From 5314f02b8f639d6382fcc7b28bd835f0022dd72e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 20:49:35 +0100 Subject: [PATCH 01/10] add `_accessors` as a place to register accessor names --- xarray/core/dataarray.py | 1 + xarray/core/dataset.py | 1 + xarray/core/datatree.py | 1 + 3 files changed, 3 insertions(+) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index d287564cfe5..1da72e394b4 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -421,6 +421,7 @@ class DataArray( _indexes: dict[Hashable, Index] _name: Hashable | None _variable: Variable + _accessors: set[str] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 74f90ce9eea..80fc7175d3e 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -709,6 +709,7 @@ class Dataset( _close: Callable[[], None] | None _indexes: dict[Hashable, Index] _variables: dict[Hashable, Variable] + _accessors: set[str] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index ee90cf7477c..932ddcb1fd2 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -224,6 +224,7 @@ class DatasetView(Dataset): "_indexes", "_variables", ) + _accessors: str[str] = set() def __init__( self, From 4dc0527871c8137267f1a3227af143b0dd135831 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 20:49:59 +0100 Subject: [PATCH 02/10] include the accessor in the new set upon registry --- xarray/core/extensions.py | 1 + 1 file changed, 1 insertion(+) 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 From bc5a25b81003d9e082bb17c24eb1c090139d883c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 20:50:25 +0100 Subject: [PATCH 03/10] include the accessor names in `__dir__` --- xarray/core/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 3a70c9ec585..703f45298fe 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -351,7 +351,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(dir(type(self))) | self._accessors | extra_attrs) def _ipython_key_completions_(self) -> list[str]: """Provide method for the key-autocompletions in IPython. From 9433b6471e5643837193720d4fabb4a53bda5291 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 20:52:55 +0100 Subject: [PATCH 04/10] actually register on `DataTree` --- xarray/core/datatree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 932ddcb1fd2..3f465b73a77 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -224,7 +224,6 @@ class DatasetView(Dataset): "_indexes", "_variables", ) - _accessors: str[str] = set() def __init__( self, @@ -456,6 +455,7 @@ class DataTree( _attrs: dict[Hashable, Any] | None _encoding: dict[Hashable, Any] | None _close: Callable[[], None] | None + _accessors: str[str] = set() __slots__ = ( "_attrs", From 2b7d0c8dbe8caa9f6b36103e020de75d9c4887e3 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 21:03:21 +0100 Subject: [PATCH 05/10] check that the registered accessors are included in `__dir__` --- xarray/tests/test_extensions.py | 4 ++++ 1 file changed, 4 insertions(+) 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" From 050d8342d16556fa3e61372b871cda217bbb0f3a Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 21:08:34 +0100 Subject: [PATCH 06/10] fix the wrong type hint --- xarray/core/datatree.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 3f465b73a77..a0a4908028d 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -455,7 +455,7 @@ class DataTree( _attrs: dict[Hashable, Any] | None _encoding: dict[Hashable, Any] | None _close: Callable[[], None] | None - _accessors: str[str] = set() + _accessors: set[str] = set() __slots__ = ( "_attrs", From 986210543534413d6dbb95c8acef354d375ab827 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 25 Jan 2025 21:09:38 +0100 Subject: [PATCH 07/10] changelog --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9b40a323f39..bcf442b493f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -73,6 +73,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 ~~~~~~~~~~~~~ From 580fd3f458743855a4ae64e99fd788ddc1d478a6 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 26 Jan 2025 15:59:19 +0100 Subject: [PATCH 08/10] mark `_accessor` as a class variable --- xarray/core/dataarray.py | 3 ++- xarray/core/dataset.py | 4 ++-- xarray/core/datatree.py | 4 ++-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 1da72e394b4..6c3cd706dd2 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, @@ -421,7 +422,7 @@ class DataArray( _indexes: dict[Hashable, Index] _name: Hashable | None _variable: Variable - _accessors: set[str] = set() + _accessors: ClassVar[set[str]] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 80fc7175d3e..89dfe47e3fc 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 @@ -709,7 +709,7 @@ class Dataset( _close: Callable[[], None] | None _indexes: dict[Hashable, Index] _variables: dict[Hashable, Variable] - _accessors: set[str] = set() + _accessors: ClassVar[set[str]] = set() __slots__ = ( "__weakref__", diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index a0a4908028d..f97cf7c91f0 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -12,7 +12,7 @@ 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 @@ -455,7 +455,7 @@ class DataTree( _attrs: dict[Hashable, Any] | None _encoding: dict[Hashable, Any] | None _close: Callable[[], None] | None - _accessors: set[str] = set() + _accessors: ClassVar[set[str]] = set() __slots__ = ( "_attrs", From bdc79840c90efd7e5bf37c9291c8e8cbebef290e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 26 Jan 2025 16:00:00 +0100 Subject: [PATCH 09/10] access `_accessors` as a class variable --- xarray/core/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 703f45298fe..77c027c3895 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -351,7 +351,7 @@ def __dir__(self) -> list[str]: for item in source if isinstance(item, str) } - return sorted(set(dir(type(self))) | self._accessors | extra_attrs) + return sorted(set(dir(type(self))) | type(self)._accessors | extra_attrs) def _ipython_key_completions_(self) -> list[str]: """Provide method for the key-autocompletions in IPython. From 336eb26ffad9e06fa18d10a66200f1299d1dfb4e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 28 Jan 2025 12:14:19 +0100 Subject: [PATCH 10/10] use a mixin to provide the extension discoverability --- xarray/core/common.py | 14 ++++++++++++-- xarray/core/dataarray.py | 8 +++++++- xarray/core/dataset.py | 2 ++ xarray/core/datatree.py | 3 ++- 4 files changed, 23 insertions(+), 4 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index 77c027c3895..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))) | type(self)._accessors | 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 6c3cd706dd2..b18a2c59903 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -39,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, @@ -282,6 +287,7 @@ class DataArray( DataWithCoords, DataArrayArithmetic, DataArrayAggregations, + AccessorMixin, ): """N-dimensional array with labeled coordinates and dimensions. diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 89dfe47e3fc..3c3691f74f2 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -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. diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index f97cf7c91f0..2465fe488ba 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -18,7 +18,7 @@ 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"], ):