diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 807baddedf9..7937a352cc6 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1035,10 +1035,12 @@ def sel( return self._from_temp_dataset(ds) def head( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "DataArray": """Return a new DataArray whose data is given by the the first `n` - values along the specified dimension(s). + values along the specified dimension(s). Default `n` = 5 See Also -------- @@ -1046,16 +1048,16 @@ def head( DataArray.tail DataArray.thin """ - - indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "head") - ds = self._to_temp_dataset().head(indexers=indexers) + ds = self._to_temp_dataset().head(indexers, **indexers_kwargs) return self._from_temp_dataset(ds) def tail( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "DataArray": """Return a new DataArray whose data is given by the the last `n` - values along the specified dimension(s). + values along the specified dimension(s). Default `n` = 5 See Also -------- @@ -1063,15 +1065,16 @@ def tail( DataArray.head DataArray.thin """ - indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "tail") - ds = self._to_temp_dataset().tail(indexers=indexers) + ds = self._to_temp_dataset().tail(indexers, **indexers_kwargs) return self._from_temp_dataset(ds) def thin( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "DataArray": """Return a new DataArray whose data is given by each `n` value - along the specified dimension(s). + along the specified dimension(s). Default `n` = 5 See Also -------- @@ -1079,8 +1082,7 @@ def thin( DataArray.head DataArray.tail """ - indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "thin") - ds = self._to_temp_dataset().thin(indexers=indexers) + ds = self._to_temp_dataset().thin(indexers, **indexers_kwargs) return self._from_temp_dataset(ds) def broadcast_like( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index d6f0da42722..1eeb5350dfe 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -2009,15 +2009,18 @@ def sel( return result._overwrite_indexes(new_indexes) def head( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "Dataset": """Returns a new dataset with the first `n` values of each array for the specified dimension(s). Parameters ---------- - indexers : dict, optional - A dict with keys matching dimensions and integer values `n`. + indexers : dict or int, default: 5 + A dict with keys matching dimensions and integer values `n` + or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. @@ -2030,20 +2033,41 @@ def head( Dataset.thin DataArray.head """ + if not indexers_kwargs: + if indexers is None: + indexers = 5 + if not isinstance(indexers, int) and not is_dict_like(indexers): + raise TypeError("indexers must be either dict-like or a single integer") + if isinstance(indexers, int): + indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "head") - indexers = {k: slice(val) for k, val in indexers.items()} - return self.isel(indexers) + for k, v in indexers.items(): + if not isinstance(v, int): + raise TypeError( + "expected integer type indexer for " + "dimension %r, found %r" % (k, type(v)) + ) + elif v < 0: + raise ValueError( + "expected positive integer as indexer " + "for dimension %r, found %s" % (k, v) + ) + indexers_slices = {k: slice(val) for k, val in indexers.items()} + return self.isel(indexers_slices) def tail( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "Dataset": """Returns a new dataset with the last `n` values of each array for the specified dimension(s). Parameters ---------- - indexers : dict, optional - A dict with keys matching dimensions and integer values `n`. + indexers : dict or int, default: 5 + A dict with keys matching dimensions and integer values `n` + or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. @@ -2056,24 +2080,44 @@ def tail( Dataset.thin DataArray.tail """ - + if not indexers_kwargs: + if indexers is None: + indexers = 5 + if not isinstance(indexers, int) and not is_dict_like(indexers): + raise TypeError("indexers must be either dict-like or a single integer") + if isinstance(indexers, int): + indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "tail") - indexers = { + for k, v in indexers.items(): + if not isinstance(v, int): + raise TypeError( + "expected integer type indexer for " + "dimension %r, found %r" % (k, type(v)) + ) + elif v < 0: + raise ValueError( + "expected positive integer as indexer " + "for dimension %r, found %s" % (k, v) + ) + indexers_slices = { k: slice(-val, None) if val != 0 else slice(val) for k, val in indexers.items() } - return self.isel(indexers) + return self.isel(indexers_slices) def thin( - self, indexers: Mapping[Hashable, Any] = None, **indexers_kwargs: Any + self, + indexers: Union[Mapping[Hashable, int], int] = None, + **indexers_kwargs: Any ) -> "Dataset": """Returns a new dataset with each array indexed along every `n`th value for the specified dimension(s) Parameters ---------- - indexers : dict, optional - A dict with keys matching dimensions and integer values `n`. + indexers : dict or int, default: 5 + A dict with keys matching dimensions and integer values `n` + or a single integer `n` applied over all dimensions. One of indexers or indexers_kwargs must be provided. **indexers_kwargs : {dim: n, ...}, optional The keyword arguments form of ``indexers``. @@ -2086,11 +2130,30 @@ def thin( Dataset.tail DataArray.thin """ + if ( + not indexers_kwargs + and not isinstance(indexers, int) + and not is_dict_like(indexers) + ): + raise TypeError("indexers must be either dict-like or a single integer") + if isinstance(indexers, int): + indexers = {dim: indexers for dim in self.dims} indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "thin") - if 0 in indexers.values(): - raise ValueError("step cannot be zero") - indexers = {k: slice(None, None, val) for k, val in indexers.items()} - return self.isel(indexers) + for k, v in indexers.items(): + if not isinstance(v, int): + raise TypeError( + "expected integer type indexer for " + "dimension %r, found %r" % (k, type(v)) + ) + elif v < 0: + raise ValueError( + "expected positive integer as indexer " + "for dimension %r, found %s" % (k, v) + ) + elif v == 0: + raise ValueError("step cannot be zero") + indexers_slices = {k: slice(None, None, val) for k, val in indexers.items()} + return self.isel(indexers_slices) def broadcast_like( self, other: Union["Dataset", "DataArray"], exclude: Iterable[Hashable] = None diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 8c01ef9a68c..78d9ace6be1 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -1005,13 +1005,48 @@ def test_isel_drop(self): def test_head(self): assert_equal(self.dv.isel(x=slice(5)), self.dv.head(x=5)) assert_equal(self.dv.isel(x=slice(0)), self.dv.head(x=0)) + assert_equal( + self.dv.isel({dim: slice(6) for dim in self.dv.dims}), self.dv.head(6) + ) + assert_equal( + self.dv.isel({dim: slice(5) for dim in self.dv.dims}), self.dv.head() + ) + with raises_regex(TypeError, "either dict-like or a single int"): + self.dv.head([3]) + with raises_regex(TypeError, "expected integer type"): + self.dv.head(x=3.1) + with raises_regex(ValueError, "expected positive int"): + self.dv.head(-3) def test_tail(self): assert_equal(self.dv.isel(x=slice(-5, None)), self.dv.tail(x=5)) assert_equal(self.dv.isel(x=slice(0)), self.dv.tail(x=0)) + assert_equal( + self.dv.isel({dim: slice(-6, None) for dim in self.dv.dims}), + self.dv.tail(6), + ) + assert_equal( + self.dv.isel({dim: slice(-5, None) for dim in self.dv.dims}), self.dv.tail() + ) + with raises_regex(TypeError, "either dict-like or a single int"): + self.dv.tail([3]) + with raises_regex(TypeError, "expected integer type"): + self.dv.tail(x=3.1) + with raises_regex(ValueError, "expected positive int"): + self.dv.tail(-3) def test_thin(self): assert_equal(self.dv.isel(x=slice(None, None, 5)), self.dv.thin(x=5)) + assert_equal( + self.dv.isel({dim: slice(None, None, 6) for dim in self.dv.dims}), + self.dv.thin(6), + ) + with raises_regex(TypeError, "either dict-like or a single int"): + self.dv.thin([3]) + with raises_regex(TypeError, "expected integer type"): + self.dv.thin(x=3.1) + with raises_regex(ValueError, "expected positive int"): + self.dv.thin(-3) with raises_regex(ValueError, "cannot be zero"): self.dv.thin(time=0) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 814fc31d734..d8401e0bd42 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -1422,6 +1422,21 @@ def test_head(self): actual = data.head(time=0) assert_equal(expected, actual) + expected = data.isel({dim: slice(6) for dim in data.dims}) + actual = data.head(6) + assert_equal(expected, actual) + + expected = data.isel({dim: slice(5) for dim in data.dims}) + actual = data.head() + assert_equal(expected, actual) + + with raises_regex(TypeError, "either dict-like or a single int"): + data.head([3]) + with raises_regex(TypeError, "expected integer type"): + data.head(dim2=3.1) + with raises_regex(ValueError, "expected positive int"): + data.head(time=-3) + def test_tail(self): data = create_test_data() @@ -1433,6 +1448,21 @@ def test_tail(self): actual = data.tail(dim1=0) assert_equal(expected, actual) + expected = data.isel({dim: slice(-6, None) for dim in data.dims}) + actual = data.tail(6) + assert_equal(expected, actual) + + expected = data.isel({dim: slice(-5, None) for dim in data.dims}) + actual = data.tail() + assert_equal(expected, actual) + + with raises_regex(TypeError, "either dict-like or a single int"): + data.tail([3]) + with raises_regex(TypeError, "expected integer type"): + data.tail(dim2=3.1) + with raises_regex(ValueError, "expected positive int"): + data.tail(time=-3) + def test_thin(self): data = create_test_data() @@ -1440,8 +1470,18 @@ def test_thin(self): actual = data.thin(time=5, dim2=6) assert_equal(expected, actual) + expected = data.isel({dim: slice(None, None, 6) for dim in data.dims}) + actual = data.thin(6) + assert_equal(expected, actual) + + with raises_regex(TypeError, "either dict-like or a single int"): + data.thin([3]) + with raises_regex(TypeError, "expected integer type"): + data.thin(dim2=3.1) with raises_regex(ValueError, "cannot be zero"): data.thin(time=0) + with raises_regex(ValueError, "expected positive int"): + data.thin(time=-3) @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_sel_fancy(self):