diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 5c1126fbc4d..c816efd083c 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -75,6 +75,8 @@ Bug fixes By `Kai Mühlbauer `_. - Fix weighted ``polyfit`` for arrays with more than two dimensions (:issue:`9972`, :pull:`9974`). By `Mattia Almansi `_. +- Include accessors in ``__dir__`` (:pull:`9985`). + By `Justus Magin `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/common.py b/xarray/core/common.py index 3a70c9ec585..51123c8f503 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -6,7 +6,7 @@ from contextlib import suppress from html import escape from textwrap import dedent -from typing import TYPE_CHECKING, Any, TypeVar, Union, overload +from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, Union, overload import numpy as np import pandas as pd @@ -262,6 +262,16 @@ def sizes(self: Any) -> Mapping[Hashable, int]: return Frozen(dict(zip(self.dims, self.shape, strict=True))) +class AccessorMixin: + """Mixin class that exposes accessors through __dir__""" + + __slots__ = () + _accessors: ClassVar[set[str]] + + def __dir__(self): + return sorted(set(super().__dir__()) | type(self)._accessors) + + class AttrAccessMixin: """Mixin class that allows getting keys with attribute access""" @@ -351,7 +361,7 @@ def __dir__(self) -> list[str]: for item in source if isinstance(item, str) } - return sorted(set(dir(type(self))) | extra_attrs) + return sorted(set(super().__dir__()) | extra_attrs) def _ipython_key_completions_(self) -> list[str]: """Provide method for the key-autocompletions in IPython. diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index d287564cfe5..b18a2c59903 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -16,6 +16,7 @@ from typing import ( TYPE_CHECKING, Any, + ClassVar, Generic, Literal, NoReturn, @@ -38,7 +39,12 @@ align, ) from xarray.core.arithmetic import DataArrayArithmetic -from xarray.core.common import AbstractArray, DataWithCoords, get_chunksizes +from xarray.core.common import ( + AbstractArray, + AccessorMixin, + DataWithCoords, + get_chunksizes, +) from xarray.core.computation import unify_chunks from xarray.core.coordinates import ( Coordinates, @@ -281,6 +287,7 @@ class DataArray( DataWithCoords, DataArrayArithmetic, DataArrayAggregations, + AccessorMixin, ): """N-dimensional array with labeled coordinates and dimensions. @@ -421,6 +428,7 @@ class DataArray( _indexes: dict[Hashable, Index] _name: Hashable | None _variable: Variable + _accessors: ClassVar[set[str]] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 74f90ce9eea..3c3691f74f2 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -24,7 +24,7 @@ from operator import methodcaller from os import PathLike from types import EllipsisType -from typing import IO, TYPE_CHECKING, Any, Generic, Literal, cast, overload +from typing import IO, TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast, overload import numpy as np from pandas.api.types import is_extension_array_dtype @@ -57,6 +57,7 @@ from xarray.core.arithmetic import DatasetArithmetic from xarray.core.array_api_compat import to_like_array from xarray.core.common import ( + AccessorMixin, DataWithCoords, _contains_datetime_like_objects, get_chunksizes, @@ -550,6 +551,7 @@ class Dataset( DataWithCoords, DatasetAggregations, DatasetArithmetic, + AccessorMixin, Mapping[Hashable, "DataArray"], ): """A multi-dimensional, in memory, array database. @@ -709,6 +711,7 @@ class Dataset( _close: Callable[[], None] | None _indexes: dict[Hashable, Index] _variables: dict[Hashable, Variable] + _accessors: ClassVar[set[str]] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index ee90cf7477c..2465fe488ba 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -12,13 +12,13 @@ Mapping, ) from html import escape -from typing import TYPE_CHECKING, Any, Literal, NoReturn, Union, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NoReturn, Union, overload from xarray.core import utils from xarray.core._aggregations import DataTreeAggregations from xarray.core._typed_ops import DataTreeOpsMixin from xarray.core.alignment import align -from xarray.core.common import TreeAttrAccessMixin, get_chunksizes +from xarray.core.common import AccessorMixin, TreeAttrAccessMixin, get_chunksizes from xarray.core.coordinates import Coordinates, DataTreeCoordinates from xarray.core.dataarray import DataArray from xarray.core.dataset import Dataset, DataVariables @@ -419,6 +419,7 @@ class DataTree( NamedNode["DataTree"], DataTreeAggregations, DataTreeOpsMixin, + AccessorMixin, TreeAttrAccessMixin, Mapping[str, "DataArray | DataTree"], ): @@ -455,6 +456,7 @@ class DataTree( _attrs: dict[Hashable, Any] | None _encoding: dict[Hashable, Any] | None _close: Callable[[], None] | None + _accessors: ClassVar[set[str]] = set() __slots__ = ( "_attrs", diff --git a/xarray/core/extensions.py b/xarray/core/extensions.py index c235fae000a..e517e656659 100644 --- a/xarray/core/extensions.py +++ b/xarray/core/extensions.py @@ -57,6 +57,7 @@ def decorator(accessor): stacklevel=2, ) setattr(cls, name, _CachedAccessor(name, accessor)) + cls._accessors.add(name) return accessor return decorator diff --git a/xarray/tests/test_extensions.py b/xarray/tests/test_extensions.py index 8a52f79198d..816693c4c19 100644 --- a/xarray/tests/test_extensions.py +++ b/xarray/tests/test_extensions.py @@ -34,6 +34,10 @@ def __init__(self, xarray_obj): def foo(self): return "bar" + assert "demo" in dir(xr.DataTree) + assert "demo" in dir(xr.Dataset) + assert "demo" in dir(xr.DataArray) + dt: xr.DataTree = xr.DataTree() assert dt.demo.foo == "bar"