Accept int value in head, thin and tail (#3298)

* Accept int value in head, thin and tail

* Fix typing

* Remove thin def val and add suggestions

* Fix typing and change raise message

Author: griverat

xarray/core/dataarray.py

@@ -1035,52 +1035,54 @@ 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
         --------
         Dataset.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
         --------
         Dataset.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
         --------
         Dataset.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(

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

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)
 

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,15 +1448,40 @@ 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()
 
         expected = data.isel(time=slice(None, None, 5), dim2=slice(None, None, 6))
         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):