Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

include accessors in __dir__ #9985

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,8 @@ Bug fixes
By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
- Fix weighted ``polyfit`` for arrays with more than two dimensions (:issue:`9972`, :pull:`9974`).
By `Mattia Almansi <https://github.com/malmans2>`_.
- Include accessors in ``__dir__`` (:pull:`9985`).
By `Justus Magin <https://github.com/keewis>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
14 changes: 12 additions & 2 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"""

Expand Down Expand Up @@ -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.
Expand Down
10 changes: 9 additions & 1 deletion xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from typing import (
TYPE_CHECKING,
Any,
ClassVar,
Generic,
Literal,
NoReturn,
Expand All @@ -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,
Expand Down Expand Up @@ -281,6 +287,7 @@ class DataArray(
DataWithCoords,
DataArrayArithmetic,
DataArrayAggregations,
AccessorMixin,
):
"""N-dimensional array with labeled coordinates and dimensions.

Expand Down Expand Up @@ -421,6 +428,7 @@ class DataArray(
_indexes: dict[Hashable, Index]
_name: Hashable | None
_variable: Variable
_accessors: ClassVar[set[str]] = set()
keewis marked this conversation as resolved.
Show resolved Hide resolved

__slots__ = (
"__weakref__",
Expand Down
5 changes: 4 additions & 1 deletion xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -550,6 +551,7 @@ class Dataset(
DataWithCoords,
DatasetAggregations,
DatasetArithmetic,
AccessorMixin,
Mapping[Hashable, "DataArray"],
):
"""A multi-dimensional, in memory, array database.
Expand Down Expand Up @@ -709,6 +711,7 @@ class Dataset(
_close: Callable[[], None] | None
_indexes: dict[Hashable, Index]
_variables: dict[Hashable, Variable]
_accessors: ClassVar[set[str]] = set()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think using a mutable default is bad practice.
But the only issue I can imagine is custom subclasses that mess with this? So maybe it's fine.

Copy link
Collaborator Author

@keewis keewis Jan 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah, the only issue would be that subclasses inherit this and try to modify this. However, if they do it would probably be better to just have them override the class variable.

(Additionally, as far as I understand the class is an instance of the metaclass, so having the usual "sentinel then switch to mutable on first access" doesn't make sense for class variables)


__slots__ = (
"__weakref__",
Expand Down
6 changes: 4 additions & 2 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -419,6 +419,7 @@ class DataTree(
NamedNode["DataTree"],
DataTreeAggregations,
DataTreeOpsMixin,
AccessorMixin,
TreeAttrAccessMixin,
Mapping[str, "DataArray | DataTree"],
):
Expand Down Expand Up @@ -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",
Expand Down
1 change: 1 addition & 0 deletions xarray/core/extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ def decorator(accessor):
stacklevel=2,
)
setattr(cls, name, _CachedAccessor(name, accessor))
cls._accessors.add(name)
return accessor

return decorator
Expand Down
4 changes: 4 additions & 0 deletions xarray/tests/test_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +37 to +39
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this PR also allow completion of the actual accessor methods, or just the accessor namespace? If so then we should test that too.

Copy link
Collaborator Author

@keewis keewis Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it does, but but how would you test that? This is functionality of the auto-completer...

Edit: to be clear, I think auto-completion of the namespace was already working before this PR

Copy link
Collaborator Author

@keewis keewis Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually, I just tried dir(ds) on main, and it included the accessor... so it might be that this PR is not necessary? Can anyone confirm that? Or did I manage to confuse myself?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it does, but but how would you test that?

Surely just

ds = xr.Dataset()
assert "foo" in dir(ds.demo)

?

I just tried dir(ds) on main, and it included the accessor... so it might be that this PR is not necessary? Can anyone confirm that? Or did I manage to confuse myself?

I just tried seeing if xr.Dataset().d would suggest xr.Dataset().dt (because that's a built-in accessor) and for me it does not.

Copy link
Collaborator Author

@keewis keewis Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dt and str are only on DataArray

Edit: And assert "foo" in dir(ds.demo) should be the same as assert "foo" in dir(DemoAccessor), no?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are any of you using a different IDE that does show the expected code-completion suggestion?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think @keewis is right, most IDEs nowadays support auto-completion via LSP servers based on (somewhat strict) static code analysers such as pyright / pylance. So I'm afraid there's nothing much we can do to support auto-completion for dynamic attributes.

IDE configs using jedi under the hood may support it, although IIUC jedi operates under different modes (Interpreter vs. Script) so I'm not sure auto-completion for dynamic attributes will always work (probably only in the Interpreter mode?).

within the IDE there's still something that has to be evaluated with regard to decorators in order for the final definition of the things that are decorated to be available to the IDE.

Decorators are just syntax sugar for functions applied on other functions or classes, so if the decorator function is properly typed (i.e., its return type is clear or can be statically inferred) any static code analyser will be able to provide auto-completions.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried to enable this for mypy here: #7117
But failed because the documentation on mypy plugins is non existent.

Maybe something similar is possible for pyright, pylance?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the insight @benbovy. Sadly, for me at least, this seems to confirm one of the reasons that I'm not a fan of dynamically registered accessors.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@max-sixty, if I remember correctly, you're also using VSCode? Do you have any additional insight here?

Sorry, just saw this.

if I understand correctly, VSCode builds its code completion on type stubs and type inference (through pylance), which means that dynamic attributes like accessors may not be possible to autocomplete.

This is also my understanding — a kernel will give better info than an IDE because it's doing something like calling dir(foo) rather than statically analyzing the code...


dt: xr.DataTree = xr.DataTree()
assert dt.demo.foo == "bar"

Expand Down
Loading