diff --git a/doc/whats-new.rst b/doc/whats-new.rst index fc5135dc598..b4dd36e9961 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -44,6 +44,7 @@ New Features - :py:meth:`xr.polyval` now supports :py:class:`Dataset` and :py:class:`DataArray` args of any shape, is faster and requires less memory. (:pull:`6548`) By `Michael Niklas `_. +- Improved overall typing. Breaking changes ~~~~~~~~~~~~~~~~ @@ -119,6 +120,9 @@ Bug fixes :pull:`6489`). By `Spencer Clark `_. - Dark themes are now properly detected in Furo-themed Sphinx documents (:issue:`6500`, :pull:`6501`). By `Kevin Paul `_. +- :py:meth:`isel` with `drop=True` works as intended with scalar :py:class:`DataArray` indexers. + (:issue:`6554`, :pull:`6579`) + By `Michael Niklas `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index fc3cbef16f8..150e5d9ca0f 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -78,7 +78,7 @@ except ImportError: iris_Cube = None - from .types import T_DataArray, T_Xarray + from .types import ErrorChoice, ErrorChoiceWithWarn, T_DataArray, T_Xarray def _infer_coords_and_dims( @@ -1171,7 +1171,7 @@ def isel( self, indexers: Mapping[Any, Any] = None, drop: bool = False, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", **indexers_kwargs: Any, ) -> DataArray: """Return a new DataArray whose data is given by integer indexing @@ -1186,7 +1186,7 @@ def isel( If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. - drop : bool, optional + drop : bool, default: False If ``drop=True``, drop coordinates variables indexed by integers instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" @@ -2335,7 +2335,7 @@ def transpose( self, *dims: Hashable, transpose_coords: bool = True, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", ) -> DataArray: """Return a new DataArray object with transposed dimensions. @@ -2386,7 +2386,7 @@ def T(self) -> DataArray: return self.transpose() def drop_vars( - self, names: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, names: Hashable | Iterable[Hashable], *, errors: ErrorChoice = "raise" ) -> DataArray: """Returns an array with dropped variables. @@ -2394,8 +2394,8 @@ def drop_vars( ---------- names : hashable or iterable of hashable Name(s) of variables to drop. - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if any of the variable + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the variable passed are not in the dataset. If 'ignore', any given names that are in the DataArray are dropped and no error is raised. @@ -2412,7 +2412,7 @@ def drop( labels: Mapping = None, dim: Hashable = None, *, - errors: str = "raise", + errors: ErrorChoice = "raise", **labels_kwargs, ) -> DataArray: """Backward compatible method based on `drop_vars` and `drop_sel` @@ -2431,7 +2431,7 @@ def drop_sel( self, labels: Mapping[Any, Any] = None, *, - errors: str = "raise", + errors: ErrorChoice = "raise", **labels_kwargs, ) -> DataArray: """Drop index labels from this DataArray. @@ -2440,8 +2440,8 @@ def drop_sel( ---------- labels : mapping of hashable to Any Index labels to drop - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the index labels passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. @@ -4589,7 +4589,7 @@ def query( queries: Mapping[Any, Any] = None, parser: str = "pandas", engine: str = None, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", **queries_kwargs: Any, ) -> DataArray: """Return a new data array indexed along the specified diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 1166e240120..f8c2223157e 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -106,7 +106,7 @@ from ..backends import AbstractDataStore, ZarrStore from .dataarray import DataArray from .merge import CoercibleMapping - from .types import T_Xarray + from .types import ErrorChoice, ErrorChoiceWithWarn, T_Xarray try: from dask.delayed import Delayed @@ -2059,7 +2059,7 @@ def chunk( return self._replace(variables) def _validate_indexers( - self, indexers: Mapping[Any, Any], missing_dims: str = "raise" + self, indexers: Mapping[Any, Any], missing_dims: ErrorChoiceWithWarn = "raise" ) -> Iterator[tuple[Hashable, int | slice | np.ndarray | Variable]]: """Here we make sure + indexer has a valid keys @@ -2164,7 +2164,7 @@ def isel( self, indexers: Mapping[Any, Any] = None, drop: bool = False, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", **indexers_kwargs: Any, ) -> Dataset: """Returns a new dataset with each array indexed along the specified @@ -2183,14 +2183,14 @@ def isel( If DataArrays are passed as indexers, xarray-style indexing will be carried out. See :ref:`indexing` for the details. One of indexers or indexers_kwargs must be provided. - drop : bool, optional + drop : bool, default: False If ``drop=True``, drop coordinates variables indexed by integers instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" What to do if dimensions that should be selected from are not present in the Dataset: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions **indexers_kwargs : {dim: indexer, ...}, optional The keyword arguments form of ``indexers``. @@ -2255,7 +2255,7 @@ def _isel_fancy( indexers: Mapping[Any, Any], *, drop: bool, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", ) -> Dataset: valid_indexers = dict(self._validate_indexers(indexers, missing_dims)) @@ -2271,6 +2271,10 @@ def _isel_fancy( } if var_indexers: new_var = var.isel(indexers=var_indexers) + # drop scalar coordinates + # https://github.com/pydata/xarray/issues/6554 + if name in self.coords and drop and new_var.ndim == 0: + continue else: new_var = var.copy(deep=False) if name not in indexes: @@ -4521,7 +4525,7 @@ def _assert_all_in_dataset( ) def drop_vars( - self, names: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, names: Hashable | Iterable[Hashable], *, errors: ErrorChoice = "raise" ) -> Dataset: """Drop variables from this dataset. @@ -4529,8 +4533,8 @@ def drop_vars( ---------- names : hashable or iterable of hashable Name(s) of variables to drop. - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if any of the variable + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the variable passed are not in the dataset. If 'ignore', any given names that are in the dataset are dropped and no error is raised. @@ -4556,7 +4560,9 @@ def drop_vars( variables, coord_names=coord_names, indexes=indexes ) - def drop(self, labels=None, dim=None, *, errors="raise", **labels_kwargs): + def drop( + self, labels=None, dim=None, *, errors: ErrorChoice = "raise", **labels_kwargs + ): """Backward compatible method based on `drop_vars` and `drop_sel` Using either `drop_vars` or `drop_sel` is encouraged @@ -4605,15 +4611,15 @@ def drop(self, labels=None, dim=None, *, errors="raise", **labels_kwargs): ) return self.drop_sel(labels, errors=errors) - def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): + def drop_sel(self, labels=None, *, errors: ErrorChoice = "raise", **labels_kwargs): """Drop index labels from this dataset. Parameters ---------- labels : mapping of hashable to Any Index labels to drop - errors : {"raise", "ignore"}, optional - If 'raise' (default), raises a ValueError error if + errors : {"raise", "ignore"}, default: "raise" + If 'raise', raises a ValueError error if any of the index labels passed are not in the dataset. If 'ignore', any given labels that are in the dataset are dropped and no error is raised. @@ -4740,7 +4746,7 @@ def drop_isel(self, indexers=None, **indexers_kwargs): return ds def drop_dims( - self, drop_dims: Hashable | Iterable[Hashable], *, errors: str = "raise" + self, drop_dims: Hashable | Iterable[Hashable], *, errors: ErrorChoice = "raise" ) -> Dataset: """Drop dimensions and associated variables from this dataset. @@ -4780,7 +4786,7 @@ def drop_dims( def transpose( self, *dims: Hashable, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", ) -> Dataset: """Return a new Dataset object with all array dimensions transposed. @@ -7714,7 +7720,7 @@ def query( queries: Mapping[Any, Any] = None, parser: str = "pandas", engine: str = None, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", **queries_kwargs: Any, ) -> Dataset: """Return a new dataset with each array indexed along the specified @@ -7747,7 +7753,7 @@ def query( Dataset: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions **queries_kwargs : {dim: query, ...}, optional diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py index e02e1f569b2..9884a756fe6 100644 --- a/xarray/core/indexes.py +++ b/xarray/core/indexes.py @@ -22,10 +22,10 @@ from . import formatting, nputils, utils from .indexing import IndexSelResult, PandasIndexingAdapter, PandasMultiIndexingAdapter -from .types import T_Index from .utils import Frozen, get_valid_numpy_dtype, is_dict_like, is_scalar if TYPE_CHECKING: + from .types import ErrorChoice, T_Index from .variable import Variable IndexVars = Dict[Any, "Variable"] @@ -1098,7 +1098,7 @@ def is_multi(self, key: Hashable) -> bool: return len(self._id_coord_names[self._coord_name_id[key]]) > 1 def get_all_coords( - self, key: Hashable, errors: str = "raise" + self, key: Hashable, errors: ErrorChoice = "raise" ) -> dict[Hashable, Variable]: """Return all coordinates having the same index. @@ -1106,7 +1106,7 @@ def get_all_coords( ---------- key : hashable Index key. - errors : {"raise", "ignore"}, optional + errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. @@ -1129,7 +1129,7 @@ def get_all_coords( return {k: self._variables[k] for k in all_coord_names} def get_all_dims( - self, key: Hashable, errors: str = "raise" + self, key: Hashable, errors: ErrorChoice = "raise" ) -> Mapping[Hashable, int]: """Return all dimensions shared by an index. @@ -1137,7 +1137,7 @@ def get_all_dims( ---------- key : hashable Index key. - errors : {"raise", "ignore"}, optional + errors : {"raise", "ignore"}, default: "raise" If "raise", raises a ValueError if `key` is not in indexes. If "ignore", an empty tuple is returned instead. diff --git a/xarray/core/types.py b/xarray/core/types.py index 74cb2fc2d46..6dbc57ce797 100644 --- a/xarray/core/types.py +++ b/xarray/core/types.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, TypeVar, Union +from typing import TYPE_CHECKING, Literal, TypeVar, Union import numpy as np @@ -33,3 +33,6 @@ DaCompatible = Union["DataArray", "Variable", "DataArrayGroupBy", "ScalarOrArray"] VarCompatible = Union["Variable", "ScalarOrArray"] GroupByIncompatible = Union["Variable", "GroupBy"] + +ErrorChoice = Literal["raise", "ignore"] +ErrorChoiceWithWarn = Literal["raise", "warn", "ignore"] diff --git a/xarray/core/utils.py b/xarray/core/utils.py index bab6476e734..aaa087a3532 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -29,6 +29,9 @@ import numpy as np import pandas as pd +if TYPE_CHECKING: + from .types import ErrorChoiceWithWarn + K = TypeVar("K") V = TypeVar("V") T = TypeVar("T") @@ -756,7 +759,9 @@ def __len__(self) -> int: def infix_dims( - dims_supplied: Collection, dims_all: Collection, missing_dims: str = "raise" + dims_supplied: Collection, + dims_all: Collection, + missing_dims: ErrorChoiceWithWarn = "raise", ) -> Iterator: """ Resolves a supplied list containing an ellipsis representing other items, to @@ -804,7 +809,7 @@ def get_temp_dimname(dims: Container[Hashable], new_dim: Hashable) -> Hashable: def drop_dims_from_indexers( indexers: Mapping[Any, Any], dims: list | Mapping[Any, int], - missing_dims: str, + missing_dims: ErrorChoiceWithWarn, ) -> Mapping[Hashable, Any]: """Depending on the setting of missing_dims, drop any dimensions from indexers that are not present in dims. @@ -850,7 +855,7 @@ def drop_dims_from_indexers( def drop_missing_dims( - supplied_dims: Collection, dims: Collection, missing_dims: str + supplied_dims: Collection, dims: Collection, missing_dims: ErrorChoiceWithWarn ) -> Collection: """Depending on the setting of missing_dims, drop any dimensions from supplied_dims that are not present in dims. diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 05c70390b46..82a567b2c2a 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -59,7 +59,7 @@ BASIC_INDEXING_TYPES = integer_types + (slice,) if TYPE_CHECKING: - from .types import T_Variable + from .types import ErrorChoiceWithWarn, T_Variable class MissingDimensionsError(ValueError): @@ -1159,7 +1159,7 @@ def _to_dense(self): def isel( self: T_Variable, indexers: Mapping[Any, Any] = None, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", **indexers_kwargs: Any, ) -> T_Variable: """Return a new array indexed along the specified dimension(s). @@ -1173,7 +1173,7 @@ def isel( What to do if dimensions that should be selected from are not present in the DataArray: - "raise": raise an exception - - "warning": raise a warning, and ignore the missing dimensions + - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions Returns @@ -1436,7 +1436,7 @@ def roll(self, shifts=None, **shifts_kwargs): def transpose( self, *dims, - missing_dims: str = "raise", + missing_dims: ErrorChoiceWithWarn = "raise", ) -> Variable: """Return a new Variable object with transposed dimensions. diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c0f7f09ff61..c1fb161fb6a 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -1175,6 +1175,25 @@ def test_isel_fancy(self): assert_array_equal(actual["var2"], expected_var2) assert_array_equal(actual["var3"], expected_var3) + # test that drop works + ds = xr.Dataset({"a": (("x",), [1, 2, 3])}, coords={"b": (("x",), [5, 6, 7])}) + + actual = ds.isel({"x": 1}, drop=False) + expected = xr.Dataset({"a": 2}, coords={"b": 6}) + assert_identical(actual, expected) + + actual = ds.isel({"x": 1}, drop=True) + expected = xr.Dataset({"a": 2}) + assert_identical(actual, expected) + + actual = ds.isel({"x": DataArray(1)}, drop=False) + expected = xr.Dataset({"a": 2}, coords={"b": 6}) + assert_identical(actual, expected) + + actual = ds.isel({"x": DataArray(1)}, drop=True) + expected = xr.Dataset({"a": 2}) + assert_identical(actual, expected) + def test_isel_dataarray(self): """Test for indexing by DataArray""" data = create_test_data()