Skip to content

Commit

Permalink
Fix Dataset/DataArray.isel with drop=True and scalar DataArray indexes (
Browse files Browse the repository at this point in the history
#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 <5635139+max-sixty@users.noreply.github.com>
  • Loading branch information
3 people authored May 10, 2022
1 parent 218e77a commit fdc3c3d
Show file tree
Hide file tree
Showing 8 changed files with 79 additions and 42 deletions.
4 changes: 4 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~~~~~~~~
Expand Down Expand Up @@ -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
~~~~~~~~~~~~~
Expand Down
24 changes: 12 additions & 12 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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`
Expand All @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down
40 changes: 23 additions & 17 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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``.
Expand Down Expand Up @@ -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))

Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
10 changes: 5 additions & 5 deletions xarray/core/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand Down
5 changes: 4 additions & 1 deletion xarray/core/types.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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"]
11 changes: 8 additions & 3 deletions xarray/core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
8 changes: 4 additions & 4 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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).
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down
Loading

0 comments on commit fdc3c3d

Please sign in to comment.