Refactor CRS-aware index

.github/workflows/ci-additional.yml

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

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,64 @@ 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 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 instead.
+    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/mixins.py

@@ -1,23 +1,17 @@
+from __future__ import annotations
+
 import abc
-import sys
 from collections.abc import Hashable
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 import pyproj
-import xarray as xr
+from xarray import DataArray, Dataset
+
+if TYPE_CHECKING:
+    from xproj.typing import CRSAwareIndex, Self
 
-try:
-    if sys.version_info >= (3, 11):
-        from typing import Self
-    else:
-        from typing_extensions import Self
-except ImportError:
-    if TYPE_CHECKING:
-        raise
-    else:
-        Self: Any = None
 
-T_Xarray_Object = TypeVar("T_Xarray_Object", xr.Dataset, xr.DataArray)
+T_Xarray_Object = TypeVar("T_Xarray_Object", Dataset, DataArray)
 
 
 class ProjAccessorMixin(abc.ABC, Generic[T_Xarray_Object]):
@@ -46,20 +40,49 @@ def _proj_set_crs(self, spatial_ref: Hashable, crs: pyproj.CRS) -> T_Xarray_Obje
 
 
 class ProjIndexMixin(abc.ABC):
-    """Mixin class that marks XProj support for an Xarray index."""
+    """Mixin class that marks XProj support for an Xarray index.
 
-    @abc.abstractmethod
-    def _proj_get_crs(self) -> pyproj.CRS | None:
-        """XProj access to the CRS of the index.
+    An :py:class:`xarray.Index` that inherits from this mixin class is
+    identified by XProj as a :term:`CRS-aware index` (note that an Xarray index
+    that simply has a ``crs`` property may also be identified as such, although
+    it may lack some XProj support).
 
-        Returns
-        -------
-        pyproj.crs.CRS or None
-            The CRS of the index or None if not (yet) defined.
+    """
 
+    @property
+    @abc.abstractmethod
+    def crs(self) -> pyproj.CRS | None:
+        """Returns the coordinate reference system (CRS) of the index as a
+        :class:`pyproj.crs.CRS` object, or ``None`` if CRS is undefined.
         """
         ...
 
+    def _proj_crs_equals(self, other: CRSAwareIndex, allow_none: bool = False) -> bool:
+        """Helper method to check if this CRS-aware index has the same CRS than the
+        other given CRS-aware index.
+
+        This method is usually called internally within the index's ``equals()``,
+        ``join()`` and ``reindex_like()`` methods.
+
+        Parameters
+        ----------
+        other : xarray.Index
+            The other CRS-aware index to compare with this index.
+        allow_none : bool, optional
+            If True, any undefined CRS is treated as the same (default: False).
+
+        """
+        # code taken from geopandas (BSD-3 Licence)
+
+        other_crs = other.crs
+
+        if allow_none:
+            if self.crs is None or other_crs is None:
+                return True
+        if not self.crs == other_crs:
+            return False
+        return True
+
     def _proj_set_crs(
         self: Self,
         spatial_ref: Hashable,