Fix Dataset/DataArray.isel with drop=True and scalar DataArray indexes (#6579)

* apply drop argument in isel_fancy

* use literal type for error handling

* add test for drop support in isel

* add isel fix to whats-new

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* correct isel unit tests

* add link to issue

* type most (all?) occurences of errors/missing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Maximilian Roos <[email protected]>

Author: 3 people

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 <https://github.com/headtr1ck>`_.
+- Improved overall typing.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
@@ -119,6 +120,9 @@ Bug fixes
   :pull:`6489`). By `Spencer Clark <https://github.com/spencerkclark>`_.
 - Dark themes are now properly detected in Furo-themed Sphinx documents (:issue:`6500`, :pull:`6501`).
   By `Kevin Paul <https://github.com/kmpaul>`_.
+- :py:meth:`isel` with `drop=True` works as intended with scalar :py:class:`DataArray` indexers.
+  (:issue:`6554`, :pull:`6579`)
+  By `Michael Niklas <https://github.com/headtr1ck>`_.
 
 Documentation
 ~~~~~~~~~~~~~

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,16 +2386,16 @@ 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.
 
         Parameters
         ----------
         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

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,16 +4525,16 @@ 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.
 
         Parameters
         ----------
         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

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,15 +1098,15 @@ 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.
 
         Parameters
         ----------
         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,15 +1129,15 @@ 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.
 
         Parameters
         ----------
         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.
 

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"]

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.

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.