speedup attribute style access and tab completion (#4742)

* speedup attribute style access and tab completion

* changes from code review, whats-new entry

Author: rhkleijn

doc/whats-new.rst

@@ -76,11 +76,12 @@ Internal Changes
   - Run the tests in parallel using pytest-xdist (:pull:`4694`).
 
   By `Justus Magin <https://github.com/keewis>`_ and `Mathias Hauser <https://github.com/mathause>`_.
-  
 - Replace all usages of ``assert x.identical(y)`` with ``assert_identical(x,  y)`` 
   for clearer error messages.
   (:pull:`4752`);
   By `Maximilian Roos <https://github.com/max-sixty>`_.
+- Speed up attribute style access (e.g. ``ds.somevar`` instead of ``ds["somevar"]``) and tab completion
+  in ipython (:issue:`4741`, :pull:`4742`). By `Richard Kleijn <https://github.com/rhkleijn>`_.
 
 .. _whats-new.0.16.2:
 

xarray/core/common.py

@@ -209,14 +209,14 @@ def __init_subclass__(cls):
             )
 
     @property
-    def _attr_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for attribute-style access"""
-        return []
+    def _attr_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for attribute-style access"""
+        yield from ()
 
     @property
-    def _item_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for key-autocompletion"""
-        return []
+    def _item_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for key-autocompletion"""
+        yield from ()
 
     def __getattr__(self, name: str) -> Any:
         if name not in {"__dict__", "__setstate__"}:
@@ -272,26 +272,26 @@ def __dir__(self) -> List[str]:
         """Provide method name lookup and completion. Only provide 'public'
         methods.
         """
-        extra_attrs = [
+        extra_attrs = set(
             item
-            for sublist in self._attr_sources
-            for item in sublist
+            for source in self._attr_sources
+            for item in source
             if isinstance(item, str)
-        ]
-        return sorted(set(dir(type(self)) + extra_attrs))
+        )
+        return sorted(set(dir(type(self))) | extra_attrs)
 
     def _ipython_key_completions_(self) -> List[str]:
         """Provide method for the key-autocompletions in IPython.
         See http://ipython.readthedocs.io/en/stable/config/integrating.html#tab-completion
         For the details.
         """
-        item_lists = [
+        items = set(
             item
-            for sublist in self._item_sources
-            for item in sublist
+            for source in self._item_sources
+            for item in source
             if isinstance(item, str)
-        ]
-        return list(set(item_lists))
+        )
+        return list(items)
 
 
 def get_squeeze_dims(

xarray/core/coordinates.py

@@ -325,29 +325,6 @@ def _ipython_key_completions_(self):
         return self._data._ipython_key_completions_()
 
 
-class LevelCoordinatesSource(Mapping[Hashable, Any]):
-    """Iterator for MultiIndex level coordinates.
-
-    Used for attribute style lookup with AttrAccessMixin. Not returned directly
-    by any public methods.
-    """
-
-    __slots__ = ("_data",)
-
-    def __init__(self, data_object: "Union[DataArray, Dataset]"):
-        self._data = data_object
-
-    def __getitem__(self, key):
-        # not necessary -- everything here can already be found in coords.
-        raise KeyError()
-
-    def __iter__(self) -> Iterator[Hashable]:
-        return iter(self._data._level_coords)
-
-    def __len__(self) -> int:
-        return len(self._data._level_coords)
-
-
 def assert_coordinate_consistent(
     obj: Union["DataArray", "Dataset"], coords: Mapping[Hashable, Variable]
 ) -> None:

xarray/core/dataarray.py

@@ -46,7 +46,6 @@
 from .common import AbstractArray, DataWithCoords
 from .coordinates import (
     DataArrayCoordinates,
-    LevelCoordinatesSource,
     assert_coordinate_consistent,
     remap_label_indexers,
 )
@@ -56,7 +55,13 @@
 from .indexing import is_fancy_indexer
 from .merge import PANDAS_TYPES, MergeError, _extract_indexes_from_coords
 from .options import OPTIONS, _get_keep_attrs
-from .utils import Default, ReprObject, _default, either_dict_or_kwargs
+from .utils import (
+    Default,
+    HybridMappingProxy,
+    ReprObject,
+    _default,
+    either_dict_or_kwargs,
+)
 from .variable import (
     IndexVariable,
     Variable,
@@ -721,18 +726,20 @@ def __delitem__(self, key: Any) -> None:
         del self.coords[key]
 
     @property
-    def _attr_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for attribute-style access"""
-        return self._item_sources + [self.attrs]
+    def _attr_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for attribute-style access"""
+        yield from self._item_sources
+        yield self.attrs
 
     @property
-    def _item_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for key-completion"""
-        return [
-            self.coords,
-            {d: self.coords[d] for d in self.dims},
-            LevelCoordinatesSource(self),
-        ]
+    def _item_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for key-completion"""
+        yield HybridMappingProxy(keys=self._coords, mapping=self.coords)
+
+        # virtual coordinates
+        # uses empty dict -- everything here can already be found in self.coords.
+        yield HybridMappingProxy(keys=self.dims, mapping={})
+        yield HybridMappingProxy(keys=self._level_coords, mapping={})
 
     def __contains__(self, key: Any) -> bool:
         return key in self.data

xarray/core/dataset.py

@@ -58,7 +58,6 @@
 )
 from .coordinates import (
     DatasetCoordinates,
-    LevelCoordinatesSource,
     assert_coordinate_consistent,
     remap_label_indexers,
 )
@@ -84,6 +83,7 @@
 from .utils import (
     Default,
     Frozen,
+    HybridMappingProxy,
     SortedKeysDict,
     _default,
     decode_numpy_dict_values,
@@ -1341,19 +1341,22 @@ def __deepcopy__(self, memo=None) -> "Dataset":
         return self.copy(deep=True)
 
     @property
-    def _attr_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for attribute-style access"""
-        return self._item_sources + [self.attrs]
+    def _attr_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for attribute-style access"""
+        yield from self._item_sources
+        yield self.attrs
 
     @property
-    def _item_sources(self) -> List[Mapping[Hashable, Any]]:
-        """List of places to look-up items for key-completion"""
-        return [
-            self.data_vars,
-            self.coords,
-            {d: self[d] for d in self.dims},
-            LevelCoordinatesSource(self),
-        ]
+    def _item_sources(self) -> Iterable[Mapping[Hashable, Any]]:
+        """Places to look-up items for key-completion"""
+        yield self.data_vars
+        yield HybridMappingProxy(keys=self._coord_names, mapping=self.coords)
+
+        # virtual coordinates
+        yield HybridMappingProxy(keys=self.dims, mapping=self)
+
+        # uses empty dict -- everything here can already be found in self.coords.
+        yield HybridMappingProxy(keys=self._level_coords, mapping={})
 
     def __contains__(self, key: object) -> bool:
         """The 'in' operator will return true or false depending on whether

xarray/core/utils.py

@@ -435,6 +435,35 @@ def FrozenDict(*args, **kwargs) -> Frozen:
     return Frozen(dict(*args, **kwargs))
 
 
+class HybridMappingProxy(Mapping[K, V]):
+    """Implements the Mapping interface. Uses the wrapped mapping for item lookup
+    and a separate wrapped keys collection for iteration.
+
+    Can be used to construct a mapping object from another dict-like object without
+    eagerly accessing its items or when a mapping object is expected but only
+    iteration over keys is actually used.
+
+    Note: HybridMappingProxy does not validate consistency of the provided `keys`
+    and `mapping`. It is the caller's responsibility to ensure that they are
+    suitable for the task at hand.
+    """
+
+    __slots__ = ("_keys", "mapping")
+
+    def __init__(self, keys: Collection[K], mapping: Mapping[K, V]):
+        self._keys = keys
+        self.mapping = mapping
+
+    def __getitem__(self, key: K) -> V:
+        return self.mapping[key]
+
+    def __iter__(self) -> Iterator[K]:
+        return iter(self._keys)
+
+    def __len__(self) -> int:
+        return len(self._keys)
+
+
 class SortedKeysDict(MutableMapping[K, V]):
     """An wrapper for dictionary-like objects that always iterates over its
     items in sorted order by key but is otherwise equivalent to the underlying