diff --git a/.github/workflows/cache-pixi-lock.yml b/.github/workflows/cache-pixi-lock.yml index e938a7664dd..da67294d62e 100644 --- a/.github/workflows/cache-pixi-lock.yml +++ b/.github/workflows/cache-pixi-lock.yml @@ -29,7 +29,7 @@ jobs: with: path: | pixi.lock - key: ${{ steps.date.outputs.date }}_${{ inputs.pixi-version }}_${{hashFiles('pixi.toml')}} + key: ${{ steps.date.outputs.date }}_${{ inputs.pixi-version }}_${{hashFiles('pixi.toml', 'pyproject.toml')}} - uses: prefix-dev/setup-pixi@v0.9.3 if: ${{ !steps.restore.outputs.cache-hit }} with: diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6db22ef1071..95f6318bf25 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -27,6 +27,12 @@ New Features brings improved alignment between h5netcdf and libnetcdf4 in the storage of complex numbers (:pull:`11068`). By `Mark Harfouche `_. +- Added typed properties for external accessor packages (hvplot, cf-xarray, + pint-xarray, rioxarray, xarray-plotly), enabling full IDE support including + autocompletion, parameter hints, and docstrings. For uninstalled packages, + ``hasattr()`` returns ``False`` to keep the namespace clean (:pull:`11079`). + By `Your Name `_. + - :py:func:`set_options` now supports an ``arithmetic_compat`` option which determines how non-index coordinates of the same name are compared for potential conflicts when performing binary operations. The default for it is ``arithmetic_compat='minimal'`` which matches the existing behaviour. diff --git a/pyproject.toml b/pyproject.toml index c8fd153dd52..917f202da1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,6 +138,7 @@ module = [ "bottleneck.*", "cartopy.*", "cf_units.*", + "cf_xarray.*", "cfgrib.*", "cftime.*", "cloudpickle.*", @@ -146,6 +147,7 @@ module = [ "fsspec.*", "h5netcdf.*", "h5py.*", + "hvplot.*", "iris.*", "mpl_toolkits.*", "nc_time_axis.*", @@ -154,14 +156,17 @@ module = [ "numcodecs.*", "opt_einsum.*", "pint.*", + "pint_xarray.*", "pooch.*", "pyarrow.*", "pydap.*", + "rioxarray.*", "scipy.*", "seaborn.*", "setuptools", "sparse.*", "toolz.*", + "xarray_plotly.*", "zarr.*", "numpy.exceptions.*", # remove once support for `numpy<2.0` has been dropped "array_api_strict.*", diff --git a/xarray/accessors.py b/xarray/accessors.py new file mode 100644 index 00000000000..515c3152474 --- /dev/null +++ b/xarray/accessors.py @@ -0,0 +1,193 @@ +""" +External accessor support for xarray. + +This module provides mixin classes with typed properties for external accessor +packages, enabling full IDE support (autocompletion, parameter hints, docstrings) +for packages like hvplot, cf-xarray, pint-xarray, rioxarray, and xarray-plotly. + +Properties are defined statically for IDE support, but raise AttributeError +for uninstalled packages (making hasattr() return False). +""" + +from __future__ import annotations + +import importlib +import importlib.util +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from cf_xarray.accessor import CFAccessor + from hvplot.xarray import hvPlotAccessor + from pint_xarray import PintDataArrayAccessor, PintDatasetAccessor + from rioxarray import RasterArray, RasterDataset + from xarray_plotly import DataArrayPlotlyAccessor, DatasetPlotlyAccessor + + from xarray.core.dataarray import DataArray + from xarray.core.dataset import Dataset + from xarray.core.datatree import DataTree + +# Registry of known external accessors +# Format: name -> (module_path, class_name, install_name, top_level_package) +DATAARRAY_ACCESSORS: dict[str, tuple[str, str, str, str]] = { + "hvplot": ("hvplot.xarray", "hvPlotAccessor", "hvplot", "hvplot"), + "cf": ("cf_xarray.accessor", "CFAccessor", "cf-xarray", "cf_xarray"), + "pint": ("pint_xarray", "PintDataArrayAccessor", "pint-xarray", "pint_xarray"), + "rio": ("rioxarray", "RasterArray", "rioxarray", "rioxarray"), + "plotly": ( + "xarray_plotly", + "DataArrayPlotlyAccessor", + "xarray-plotly", + "xarray_plotly", + ), +} + +DATASET_ACCESSORS: dict[str, tuple[str, str, str, str]] = { + "hvplot": ("hvplot.xarray", "hvPlotAccessor", "hvplot", "hvplot"), + "cf": ("cf_xarray.accessor", "CFAccessor", "cf-xarray", "cf_xarray"), + "pint": ("pint_xarray", "PintDatasetAccessor", "pint-xarray", "pint_xarray"), + "rio": ("rioxarray", "RasterDataset", "rioxarray", "rioxarray"), + "plotly": ( + "xarray_plotly", + "DatasetPlotlyAccessor", + "xarray-plotly", + "xarray_plotly", + ), +} + +DATATREE_ACCESSORS: dict[str, tuple[str, str, str, str]] = { + "hvplot": ("hvplot.xarray", "hvPlotAccessor", "hvplot", "hvplot"), + "cf": ("cf_xarray.accessor", "CFAccessor", "cf-xarray", "cf_xarray"), +} + +# Cache for package availability checks +_package_available_cache: dict[str, bool] = {} + + +def _is_package_available(package_name: str) -> bool: + """Check if a package is available without importing it.""" + if package_name not in _package_available_cache: + _package_available_cache[package_name] = ( + importlib.util.find_spec(package_name) is not None + ) + return _package_available_cache[package_name] + + +def _get_external_accessor( + name: str, + obj: DataArray | Dataset | DataTree, + accessor_registry: dict[str, tuple[str, str, str, str]], +) -> Any: + """Get an external accessor instance, raising AttributeError if not installed.""" + package, cls_name, install_name, top_pkg = accessor_registry[name] + + if not _is_package_available(top_pkg): + raise AttributeError( + f"'{type(obj).__name__}' object has no attribute '{name}'. " + f"Install with: pip install {install_name}" + ) + + # Check cache + try: + cache = obj._cache + except AttributeError: + cache = obj._cache = {} + + cache_key = f"_external_{name}" + if cache_key in cache: + return cache[cache_key] + + # Import and instantiate the accessor + try: + module = importlib.import_module(package) + except ImportError as err: + raise AttributeError( + f"'{type(obj).__name__}' object has no attribute '{name}'. " + f"Install with: pip install {install_name}" + ) from err + + try: + accessor_cls = getattr(module, cls_name) + accessor = accessor_cls(obj) + except AttributeError as err: + raise RuntimeError(f"Error initializing {name!r} accessor.") from err + + cache[cache_key] = accessor + return accessor + + +class DataArrayExternalAccessorMixin: + """Mixin providing typed external accessor properties for DataArray.""" + + __slots__ = () + + @property + def hvplot(self) -> hvPlotAccessor: + """hvPlot accessor for interactive plotting. Requires: ``pip install hvplot``""" + return _get_external_accessor("hvplot", self, DATAARRAY_ACCESSORS) # type: ignore[arg-type] + + @property + def cf(self) -> CFAccessor: + """CF conventions accessor. Requires: ``pip install cf-xarray``""" + return _get_external_accessor("cf", self, DATAARRAY_ACCESSORS) # type: ignore[arg-type] + + @property + def pint(self) -> PintDataArrayAccessor: + """Pint unit accessor. Requires: ``pip install pint-xarray``""" + return _get_external_accessor("pint", self, DATAARRAY_ACCESSORS) # type: ignore[arg-type] + + @property + def rio(self) -> RasterArray: + """Rasterio accessor for geospatial data. Requires: ``pip install rioxarray``""" + return _get_external_accessor("rio", self, DATAARRAY_ACCESSORS) # type: ignore[arg-type] + + @property + def plotly(self) -> DataArrayPlotlyAccessor: + """Plotly accessor. Requires: ``pip install xarray-plotly``""" + return _get_external_accessor("plotly", self, DATAARRAY_ACCESSORS) # type: ignore[arg-type] + + +class DatasetExternalAccessorMixin: + """Mixin providing typed external accessor properties for Dataset.""" + + __slots__ = () + + @property + def hvplot(self) -> hvPlotAccessor: + """hvPlot accessor for interactive plotting. Requires: ``pip install hvplot``""" + return _get_external_accessor("hvplot", self, DATASET_ACCESSORS) # type: ignore[arg-type] + + @property + def cf(self) -> CFAccessor: + """CF conventions accessor. Requires: ``pip install cf-xarray``""" + return _get_external_accessor("cf", self, DATASET_ACCESSORS) # type: ignore[arg-type] + + @property + def pint(self) -> PintDatasetAccessor: + """Pint unit accessor. Requires: ``pip install pint-xarray``""" + return _get_external_accessor("pint", self, DATASET_ACCESSORS) # type: ignore[arg-type] + + @property + def rio(self) -> RasterDataset: + """Rasterio accessor for geospatial data. Requires: ``pip install rioxarray``""" + return _get_external_accessor("rio", self, DATASET_ACCESSORS) # type: ignore[arg-type] + + @property + def plotly(self) -> DatasetPlotlyAccessor: + """Plotly accessor. Requires: ``pip install xarray-plotly``""" + return _get_external_accessor("plotly", self, DATASET_ACCESSORS) # type: ignore[arg-type] + + +class DataTreeExternalAccessorMixin: + """Mixin providing typed external accessor properties for DataTree.""" + + __slots__ = () + + @property + def hvplot(self) -> hvPlotAccessor: + """hvPlot accessor for interactive plotting. Requires: ``pip install hvplot``""" + return _get_external_accessor("hvplot", self, DATATREE_ACCESSORS) # type: ignore[arg-type] + + @property + def cf(self) -> CFAccessor: + """CF conventions accessor. Requires: ``pip install cf-xarray``""" + return _get_external_accessor("cf", self, DATATREE_ACCESSORS) # type: ignore[arg-type] diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 706c35e5459..b4625efc43d 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -19,6 +19,7 @@ import numpy as np import pandas as pd +from xarray.accessors import DataArrayExternalAccessorMixin from xarray.coding.calendar_ops import convert_calendar, interp_calendar from xarray.coding.cftimeindex import CFTimeIndex from xarray.computation import computation, ops @@ -259,6 +260,7 @@ class DataArray( DataWithCoords, DataArrayArithmetic, DataArrayAggregations, + DataArrayExternalAccessorMixin, ): """N-dimensional array with labeled coordinates and dimensions. diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index e15f1077639..e1f01c2e1f0 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -29,6 +29,7 @@ import numpy as np import pandas as pd +from xarray.accessors import DatasetExternalAccessorMixin from xarray.coding.calendar_ops import convert_calendar, interp_calendar from xarray.coding.cftimeindex import CFTimeIndex, _parse_array_of_cftime_strings from xarray.compat.array_api_compat import to_like_array @@ -197,6 +198,7 @@ class Dataset( DataWithCoords, DatasetAggregations, DatasetArithmetic, + DatasetExternalAccessorMixin, Mapping[Hashable, "DataArray"], ): """A multi-dimensional, in memory, array database. diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index e079332780c..d46fd236feb 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -28,6 +28,7 @@ overload, ) +from xarray.accessors import DataTreeExternalAccessorMixin from xarray.core import utils from xarray.core._aggregations import DataTreeAggregations from xarray.core._typed_ops import DataTreeOpsMixin @@ -464,6 +465,7 @@ class DataTree( DataTreeAggregations, DataTreeOpsMixin, TreeAttrAccessMixin, + DataTreeExternalAccessorMixin, Mapping[str, "DataArray | DataTree"], ): """ diff --git a/xarray/core/extensions.py b/xarray/core/extensions.py index c235fae000a..51a7917b399 100644 --- a/xarray/core/extensions.py +++ b/xarray/core/extensions.py @@ -50,11 +50,28 @@ def __get__(self, obj, cls): def _register_accessor(name, cls): def decorator(accessor): if hasattr(cls, name): + # Skip registration for known external accessors - xarray provides + # typed properties that load them directly for IDE support + from xarray.accessors import ( + DATAARRAY_ACCESSORS, + DATASET_ACCESSORS, + DATATREE_ACCESSORS, + ) + + known_external = ( + set(DATAARRAY_ACCESSORS) + | set(DATASET_ACCESSORS) + | set(DATATREE_ACCESSORS) + ) + if name in known_external: + # Don't overwrite - our typed property handles this accessor + return accessor + warnings.warn( f"registration of accessor {accessor!r} under name {name!r} for type {cls!r} is " "overriding a preexisting attribute with the same name.", AccessorRegistrationWarning, - stacklevel=2, + stacklevel=3, ) setattr(cls, name, _CachedAccessor(name, accessor)) return accessor diff --git a/xarray/tests/test_extensions.py b/xarray/tests/test_extensions.py index 8a52f79198d..8c614b82b0b 100644 --- a/xarray/tests/test_extensions.py +++ b/xarray/tests/test_extensions.py @@ -5,6 +5,12 @@ import pytest import xarray as xr +from xarray.accessors import ( + DATAARRAY_ACCESSORS, + DATASET_ACCESSORS, + DATATREE_ACCESSORS, + _is_package_available, +) from xarray.core.extensions import register_datatree_accessor from xarray.tests import assert_identical @@ -93,3 +99,74 @@ def __init__(self, xarray_obj): with pytest.raises(RuntimeError, match=r"error initializing"): _ = xr.Dataset().stupid_accessor + + +class TestExternalAccessors: + """Tests for typed external accessor properties.""" + + def test_hasattr_false_for_uninstalled(self) -> None: + """hasattr returns False for accessors whose packages are not installed.""" + da = xr.DataArray([1, 2, 3]) + ds = xr.Dataset({"a": da}) + dt = xr.DataTree(ds) + + for name, (_, _, _, top_pkg) in DATAARRAY_ACCESSORS.items(): + if not _is_package_available(top_pkg): + assert not hasattr(da, name), f"hasattr should be False for {name}" + + for name, (_, _, _, top_pkg) in DATASET_ACCESSORS.items(): + if not _is_package_available(top_pkg): + assert not hasattr(ds, name), f"hasattr should be False for {name}" + + for name, (_, _, _, top_pkg) in DATATREE_ACCESSORS.items(): + if not _is_package_available(top_pkg): + assert not hasattr(dt, name), f"hasattr should be False for {name}" + + def test_hasattr_true_for_installed(self) -> None: + """hasattr returns True for accessors whose packages are installed.""" + da = xr.DataArray([1, 2, 3]) + ds = xr.Dataset({"a": da}) + + for name, (_, _, _, top_pkg) in DATAARRAY_ACCESSORS.items(): + if _is_package_available(top_pkg): + assert hasattr(da, name), f"hasattr should be True for {name}" + + for name, (_, _, _, top_pkg) in DATASET_ACCESSORS.items(): + if _is_package_available(top_pkg): + assert hasattr(ds, name), f"hasattr should be True for {name}" + + def test_attribute_error_for_uninstalled(self) -> None: + """Accessing uninstalled accessor raises AttributeError.""" + da = xr.DataArray([1, 2, 3]) + + for name, (_, _, _, top_pkg) in DATAARRAY_ACCESSORS.items(): + if not _is_package_available(top_pkg): + with pytest.raises(AttributeError): + getattr(da, name) + break # Only need to test one + + def test_external_accessor_no_overwrite(self) -> None: + """Known external accessors don't overwrite typed properties.""" + # The property should remain a property, not get replaced by _CachedAccessor + for name in DATAARRAY_ACCESSORS: + attr = getattr(xr.DataArray, name) + assert isinstance(attr, property), f"{name} should remain a property" + + for name in DATASET_ACCESSORS: + attr = getattr(xr.Dataset, name) + assert isinstance(attr, property), f"{name} should remain a property" + + for name in DATATREE_ACCESSORS: + attr = getattr(xr.DataTree, name) + assert isinstance(attr, property), f"{name} should remain a property" + + def test_accessor_caching(self) -> None: + """Accessor instances are cached on the object.""" + da = xr.DataArray([1, 2, 3]) + + for name, (_, _, _, top_pkg) in DATAARRAY_ACCESSORS.items(): + if _is_package_available(top_pkg): + accessor1 = getattr(da, name) + accessor2 = getattr(da, name) + assert accessor1 is accessor2, f"{name} accessor should be cached" + break # Only need to test one installed accessor