Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 10 commits into from
May 10, 2022
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
headtr1ck marked this conversation as resolved.
Show resolved Hide resolved
- "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"]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great! We could find / replace for these at some point

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even though it is scope creep, I have typed any occurences that I have found in this PR.
Mypy is not complaining locally about these changes.

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