Wrap an array to make vectorized indexing lazy.
| 817 | |
| 818 | |
| 819 | class LazilyVectorizedIndexedArray(ExplicitlyIndexedNDArrayMixin): |
| 820 | """Wrap an array to make vectorized indexing lazy.""" |
| 821 | |
| 822 | __slots__ = ("array", "key") |
| 823 | |
| 824 | def __init__(self, array: duckarray[Any, Any], key: ExplicitIndexer): |
| 825 | """ |
| 826 | Parameters |
| 827 | ---------- |
| 828 | array : array_like |
| 829 | Array like object to index. |
| 830 | key : VectorizedIndexer |
| 831 | """ |
| 832 | if isinstance(key, BasicIndexer | OuterIndexer): |
| 833 | self.key = _outer_to_vectorized_indexer(key, array.shape) |
| 834 | elif isinstance(key, VectorizedIndexer): |
| 835 | self.key = _arrayize_vectorized_indexer(key, array.shape) |
| 836 | self.array = as_indexable(array) |
| 837 | |
| 838 | @property |
| 839 | def shape(self) -> _Shape: |
| 840 | return np.broadcast(*self.key.tuple).shape |
| 841 | |
| 842 | def get_duck_array(self): |
| 843 | from xarray.backends.common import BackendArray |
| 844 | |
| 845 | if isinstance(self.array, BackendArray): |
| 846 | array = self.array[self.key] |
| 847 | else: |
| 848 | array = apply_indexer(self.array, self.key) |
| 849 | if isinstance(array, ExplicitlyIndexed): |
| 850 | array = array.get_duck_array() |
| 851 | return _wrap_numpy_scalars(array) |
| 852 | |
| 853 | async def async_get_duck_array(self): |
| 854 | from xarray.backends.common import BackendArray |
| 855 | |
| 856 | if isinstance(self.array, BackendArray): |
| 857 | array = await self.array.async_getitem(self.key) |
| 858 | else: |
| 859 | array = apply_indexer(self.array, self.key) |
| 860 | if isinstance(array, ExplicitlyIndexed): |
| 861 | array = await array.async_get_duck_array() |
| 862 | return _wrap_numpy_scalars(array) |
| 863 | |
| 864 | def _updated_key(self, new_key: ExplicitIndexer): |
| 865 | return _combine_indexers(self.key, self.shape, new_key) |
| 866 | |
| 867 | def _oindex_get(self, indexer: OuterIndexer): |
| 868 | return type(self)(self.array, self._updated_key(indexer)) |
| 869 | |
| 870 | def _vindex_get(self, indexer: VectorizedIndexer): |
| 871 | return type(self)(self.array, self._updated_key(indexer)) |
| 872 | |
| 873 | def __getitem__(self, indexer: ExplicitIndexer): |
| 874 | self._check_and_raise_if_non_basic_indexer(indexer) |
| 875 | # If the indexed array becomes a scalar, return LazilyIndexedArray |
| 876 | if all(isinstance(ind, integer_types) for ind in indexer.tuple): |
no outgoing calls
no test coverage detected
searching dependent graphs…