Refactor CRS-aware index (#27)

* refactor CRS-aware index

ProjIndexMixin class:

- replace the ``_proj_get_crs`` abstract method by the ``crs``
abstract (public) property.

- add ``_proj_crs_equals`` helper method

Xproj now identifies an Xarray index as CRS-aware if it has a ``crs``
property that returns either a ``pyproj.CRS`` object or ``None``. A
CRS-aware index may have an undefined CRS (``None``).

Add ``CRSAwareIndex`` protocol class.

* add a couple of CRS utility functions

Copied and adapted from xvec.

* CI: mypy install additional types

* doc: add crs property in mixin class API

* docstring tweaks

* misc. typing fixes

* ci: temp fix (xarray upstream)

Author: benbovy

.github/workflows/ci-additional.yml

@@ -45,4 +45,4 @@ jobs:
 
       - name: Run mypy
         run: |
-          python -m mypy
+          python -m mypy --install-types --non-interactive

.github/workflows/test.yml

@@ -37,6 +37,8 @@ jobs:
       - name: Install xproj
         run: |
           python -m pip install -e .[test]
+          # TODO: remove in next xarray release (https://github.com/pydata/xarray/pull/10413)
+          python -m pip install typing_extensions
           python -m pip install pytest-cov
 
       - name: Run tests

docs/_templates/autosummary/mixin.rst

@@ -8,7 +8,7 @@
 
    .. autosummary::
    {% for item in members %}
-     {% if item.startswith('_proj') %}
+     {% if item.startswith('_proj') or item == "crs" %}
       ~{{ name }}.{{ item }}
      {% endif %}
    {%- endfor %}

docs/api-hidden.rst

@@ -12,7 +12,8 @@
    :toctree: _api_generated/
 
    ProjAccessorMixin._proj_set_crs
-   ProjIndexMixin._proj_get_crs
+   ProjIndexMixin.crs
+   ProjIndexMixin._proj_crs_equals
    ProjIndexMixin._proj_set_crs
    ProjIndexMixin._proj_to_crs
    format_compact_cf

docs/api.rst

@@ -74,6 +74,15 @@ To enable it, be sure to import ``xproj`` after ``xarray``:
 
 .. currentmodule:: xproj
 
+CRS utility functions
+---------------------
+
+.. autosummary::
+   :toctree: _api_generated/
+
+   format_crs
+   get_common_crs
+
 3rd-party Xarray extensions
 ---------------------------
 

docs/integration.md

@@ -126,7 +126,8 @@ class GeoIndex(xr.indexes.PandasIndex, xproj.ProjIndexMixin):
 
         return super().sel(*args, **kwargs)
 
-    def _proj_get_crs(self):
+    @property
+    def crs(self):
         return self._crs
 
     def _proj_set_crs(self, spatial_ref, crs):
@@ -184,8 +185,9 @@ CRS-aware! (just a warning is emitted below).
 ds_geo_wgs84.sel(lat=70)
 ```
 
-Since ``GeoIndex`` also implements the ``_proj_get_crs`` method it is possible
-to get the CRS from the "lat" coordinate like so:
+``GeoIndex`` has a ``crs`` property (as required by
+{class}`~xproj.ProjIndexMixin`), which is possible to access also via the
+``proj`` accessor like so:
 
 ```{code-cell} ipython3
 ds_geo_wgs84.proj("lat").crs

xproj/__init__.py

@@ -2,7 +2,7 @@
 
 from .accessor import ProjAccessor as _ProjAccessor  # noqa: F401
 from .accessor import register_accessor
-from .crs_utils import format_compact_cf, format_full_cf_gdal
+from .crs_utils import format_compact_cf, format_crs, format_full_cf_gdal, get_common_crs
 from .index import CRSIndex  # noqa: F401
 from .mixins import ProjAccessorMixin, ProjIndexMixin
 
@@ -12,7 +12,9 @@
     "ProjAccessorMixin",
     "ProjIndexMixin",
     "format_compact_cf",
+    "format_crs",
     "format_full_cf_gdal",
+    "get_common_crs",
     "register_accessor",
 ]
 

xproj/accessor.py

@@ -10,6 +10,7 @@
 from xproj.crs_utils import format_compact_cf
 from xproj.index import CRSIndex
 from xproj.mixins import ProjIndexMixin
+from xproj.typing import CRSAwareIndex
 from xproj.utils import Frozen, FrozenDict
 
 
@@ -94,21 +95,31 @@ class CRSProxy:
 
     _obj: xr.Dataset | xr.DataArray
     _crs_coord_name: Hashable
-    _crs: pyproj.CRS
+    _crs: pyproj.CRS | None
 
-    def __init__(self, obj: xr.Dataset | xr.DataArray, coord_name: Hashable, crs: pyproj.CRS):
+    def __init__(
+        self, obj: xr.Dataset | xr.DataArray, coord_name: Hashable, crs: pyproj.CRS | None
+    ):
         self._obj = obj
         self._crs_coord_name = coord_name
         self._crs = crs
 
     @property
-    def crs(self) -> pyproj.CRS:
-        """Return the coordinate reference system as a :class:`pyproj.CRS` object."""
+    def crs(self) -> pyproj.CRS | None:
+        """Return the coordinate reference system as a :class:`pyproj.CRS` object, or
+        ``None`` if the CRS is undefined.
+        """
         return self._crs
 
 
 def is_crs_aware(index: xr.Index) -> bool:
-    return isinstance(index, ProjIndexMixin) or hasattr(index, "_proj_get_crs")
+    if isinstance(index, ProjIndexMixin):
+        return True
+    if hasattr(index, "crs"):
+        crs = getattr(index, "crs")
+        if isinstance(crs, pyproj.CRS) or crs is None:
+            return True
+    return False
 
 
 @xr.register_dataset_accessor("proj")
@@ -118,7 +129,7 @@ class ProjAccessor:
 
     _obj: xr.Dataset | xr.DataArray
     _crs_indexes: dict[Hashable, CRSIndex] | None
-    _crs_aware_indexes: dict[Hashable, xr.Index] | None
+    _crs_aware_indexes: dict[Hashable, CRSAwareIndex] | None
     _crs: pyproj.CRS | None | Literal[False]
 
     def __init__(self, obj: xr.Dataset | xr.DataArray):
@@ -138,7 +149,7 @@ def _cache_all_crs_indexes(self):
                 self._crs_indexes[name] = idx
             elif is_crs_aware(idx):
                 for name in vars:
-                    self._crs_aware_indexes[name] = idx
+                    self._crs_aware_indexes[name] = cast(CRSAwareIndex, idx)
 
     @property
     def crs_indexes(self) -> Frozen[Hashable, CRSIndex]:
@@ -154,13 +165,12 @@ def crs_indexes(self) -> Frozen[Hashable, CRSIndex]:
         return FrozenDict(self._crs_indexes)
 
     @property
-    def crs_aware_indexes(self) -> Frozen[Hashable, xr.Index]:
+    def crs_aware_indexes(self) -> Frozen[Hashable, CRSAwareIndex]:
         """Return an immutable dictionary of coordinate names as keys and
         xarray Index objects that are CRS-aware.
 
         A :term:`CRS-aware index` is an :py:class:`xarray.Index` object that
-        must at least implements a method like
-        :py:meth:`~xproj.ProjIndexMixin._proj_get_crs`.
+        must at least implement a property like :py:meth:`~xproj.ProjIndexMixin.crs`.
 
         """
         if self._crs_aware_indexes is None:
@@ -205,10 +215,10 @@ def __call__(self, coord_name: Hashable):
             A proxy accessor for a single CRS.
 
         """
-        crs: pyproj.CRS
+        crs: pyproj.CRS | None
 
         if coord_name in self.crs_aware_indexes:
-            crs = self.crs_aware_indexes[coord_name]._proj_get_crs()  # type: ignore
+            crs = self.crs_aware_indexes[coord_name].crs
         else:
             crs = self._get_crs_index(coord_name).crs
 
@@ -236,7 +246,7 @@ def crs(self) -> pyproj.CRS | None:
         if self._crs is False:
             all_crs = {name: idx.crs for name, idx in self.crs_indexes.items()}
             for name, idx in self.crs_aware_indexes.items():
-                crs = idx._proj_get_crs()  # type: ignore
+                crs = idx.crs
                 if crs is not None:
                     all_crs[name] = crs
 
@@ -387,7 +397,7 @@ def map_crs(
                     )
                     continue
 
-                index_crs = index._proj_get_crs()  # type: ignore
+                index_crs = cast(CRSAwareIndex, index).crs
 
                 if not allow_override:
                     if index_crs is not None and index_crs != crs:

xproj/crs_utils.py

@@ -1,9 +1,31 @@
-from typing import Any
+import warnings
+from collections.abc import Sequence
+from typing import Any, Literal
 
-import pyproj
+from pyproj import CRS
 
 
-def format_compact_cf(crs: pyproj.CRS) -> dict[str, Any]:
+def format_crs(crs: CRS | None, max_width: int = 50) -> str:
+    """Format CRS as a string.
+
+    Parameters
+    ----------
+    crs : pyproj.crs.CRS
+        The input CRS object to format.
+    max_width : int, optional
+        Maximum number of characters beyond which the formatted CRS
+        will be truncated (default: 50).
+
+    """
+    if crs is not None:
+        srs = crs.to_string()
+    else:
+        srs = "None"
+
+    return srs if len(srs) <= max_width else " ".join([srs[:max_width], "..."])
+
+
+def format_compact_cf(crs: CRS) -> dict[str, Any]:
     """Format CRS as a dictionary for minimal compatibility with
     CF conventions.
 
@@ -30,7 +52,7 @@ def format_compact_cf(crs: pyproj.CRS) -> dict[str, Any]:
     return {"crs_wkt": crs.to_wkt()}
 
 
-def format_full_cf_gdal(crs: pyproj.CRS) -> dict[str, Any]:
+def format_full_cf_gdal(crs: CRS) -> dict[str, Any]:
     """Format CRS as a dictionary for full compatibility with
     CF conventions and GDAL.
 
@@ -61,3 +83,65 @@ def format_full_cf_gdal(crs: pyproj.CRS) -> dict[str, Any]:
     output = crs.to_cf()
     output["spatial_ref"] = crs.to_wkt()
     return output
+
+
+def get_common_crs(
+    crs_objs: Sequence[CRS | None] | set[CRS | None],
+    on_undefined_crs: Literal["raise", "warn", "ignore"] = "warn",
+    stacklevel: int = 3,
+) -> CRS | None:
+    """Try getting a common, unique CRS from an input sequence of (possibly
+    undefined) CRS objects.
+
+    Parameters
+    ----------
+    crs_objs : sequence or set
+        Sequence of either :py:class:`pyproj.CRS` objects or ``None``
+        (undefined CRS).
+    on_undefined_crs : {"raise", "warn", "ignore"}, optional
+        If "raise", raises a ValueError if a non-null CRS is found but
+        one or more inputs have undefined CRS. If "warn" (default), emits a
+        UserWarning instead. If "ignore", do nothing.
+    stacklevel : int, optional
+        Stack level value used for the emitted warning (default: 3).
+
+    Returns
+    -------
+    pyproj.crs.CRS or None
+        The common (possibly undefined) CRS.
+
+    Raises
+    ------
+    ValueError
+        If multiple conflicting CRS objects are found.
+
+    Warns
+    -----
+    UserWarning
+        If a common, unique CRS is found but one or more of the
+        inputs have undefined CRS.
+
+    """
+    # code taken from geopandas (BSD-3 Licence)
+
+    crs_objs = set(crs_objs)
+
+    crs_not_none = [crs for crs in crs_objs if crs is not None]
+    names = [crs.name for crs in crs_not_none]
+
+    if len(crs_not_none) == 0:
+        return None
+    if len(crs_not_none) == 1:
+        if len(crs_objs) != 1:
+            if on_undefined_crs == "raise":
+                raise ValueError("one or more inputs have undefined CRS.")
+            elif on_undefined_crs == "warn":
+                warnings.warn(  # noqa: B028
+                    "CRS is undefined for some of the inputs. "
+                    f"Setting output's CRS as {names[0]} "
+                    "(the single non-null CRS provided).",
+                    stacklevel=stacklevel,
+                )
+        return crs_not_none[0]
+
+    raise ValueError(f"cannot determine common CRS from inputs CRSes {names}. ")

xproj/index.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
-from _collections_abc import Mapping
+from collections.abc import Hashable, Mapping
 from typing import Any
 
 import pyproj
 import xarray as xr
+from pyproj.exceptions import CRSError
 from xarray.indexes import Index
 
 
@@ -69,15 +70,15 @@ def from_variables(
                 crs = pyproj.CRS.from_user_input(options["crs"])
             else:
                 crs = pyproj.CRS.from_cf(var.attrs)
-        except pyproj.crs.CRSError:
+        except CRSError:
             raise ValueError(
                 f"CRS could not be constructed from attrs on provided variable {varname!r}"
                 f"Either add appropriate attributes to {varname!r} or pass a `crs` kwarg."
             )
 
         return cls(crs)
 
-    def equals(self, other: Index) -> bool:
+    def equals(self, other: Index, *, exclude: frozenset[Hashable] | None = None) -> bool:
         if not isinstance(other, CRSIndex):
             return False
         if not self.crs == other.crs: