Use Self rather than concrete types, remove casts (#8216)

* Use `Self` rather than concrete types, remove `cast`s

This should also allow for subtyping

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

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

* Undo one `Self`

* Unused ignore

* Add check for redundant self annotations

* And `DataWithCoords`

* And `DataArray` & `Dataset`

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

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

* And `Variable`

* Update xarray/core/dataarray.py

Co-authored-by: Michael Niklas  <[email protected]>

* Update xarray/core/dataarray.py

Co-authored-by: Michael Niklas  <[email protected]>

* Update xarray/core/dataarray.py

Co-authored-by: Michael Niklas  <[email protected]>

* Clean-ups — `other`, casts, obsolete comments

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

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

* another one

---------

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

Author: 3 people

pyproject.toml

@@ -72,6 +72,7 @@ source = ["xarray"]
 exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"]
 
 [tool.mypy]
+enable_error_code = "redundant-self"
 exclude = 'xarray/util/generate_.*\.py'
 files = "xarray"
 show_error_codes = true

xarray/core/accessor_str.py

@@ -2386,7 +2386,7 @@ def _partitioner(
 
         # _apply breaks on an empty array in this case
         if not self._obj.size:
-            return self._obj.copy().expand_dims({dim: 0}, axis=-1)  # type: ignore[return-value]
+            return self._obj.copy().expand_dims({dim: 0}, axis=-1)
 
         arrfunc = lambda x, isep: np.array(func(x, isep), dtype=self._obj.dtype)
 

xarray/core/common.py

@@ -45,6 +45,7 @@
         DatetimeLike,
         DTypeLikeSave,
         ScalarOrArray,
+        Self,
         SideOptions,
         T_Chunks,
         T_DataWithCoords,
@@ -381,11 +382,11 @@ class DataWithCoords(AttrAccessMixin):
     __slots__ = ("_close",)
 
     def squeeze(
-        self: T_DataWithCoords,
+        self,
         dim: Hashable | Iterable[Hashable] | None = None,
         drop: bool = False,
         axis: int | Iterable[int] | None = None,
-    ) -> T_DataWithCoords:
+    ) -> Self:
         """Return a new object with squeezed data.
 
         Parameters
@@ -414,12 +415,12 @@ def squeeze(
         return self.isel(drop=drop, **{d: 0 for d in dims})
 
     def clip(
-        self: T_DataWithCoords,
+        self,
         min: ScalarOrArray | None = None,
         max: ScalarOrArray | None = None,
         *,
         keep_attrs: bool | None = None,
-    ) -> T_DataWithCoords:
+    ) -> Self:
         """
         Return an array whose values are limited to ``[min, max]``.
         At least one of max or min must be given.
@@ -472,10 +473,10 @@ def _calc_assign_results(
         return {k: v(self) if callable(v) else v for k, v in kwargs.items()}
 
     def assign_coords(
-        self: T_DataWithCoords,
+        self,
         coords: Mapping[Any, Any] | None = None,
         **coords_kwargs: Any,
-    ) -> T_DataWithCoords:
+    ) -> Self:
         """Assign new coordinates to this object.
 
         Returns a new object with all the original data in addition to the new
@@ -620,9 +621,7 @@ def assign_coords(
         data.coords.update(results)
         return data
 
-    def assign_attrs(
-        self: T_DataWithCoords, *args: Any, **kwargs: Any
-    ) -> T_DataWithCoords:
+    def assign_attrs(self, *args: Any, **kwargs: Any) -> Self:
         """Assign new attrs to this object.
 
         Returns a new object equivalent to ``self.attrs.update(*args, **kwargs)``.
@@ -1061,9 +1060,7 @@ def _resample(
             restore_coord_dims=restore_coord_dims,
         )
 
-    def where(
-        self: T_DataWithCoords, cond: Any, other: Any = dtypes.NA, drop: bool = False
-    ) -> T_DataWithCoords:
+    def where(self, cond: Any, other: Any = dtypes.NA, drop: bool = False) -> Self:
         """Filter elements from this object according to a condition.
 
         Returns elements from 'DataArray', where 'cond' is True,
@@ -1208,9 +1205,7 @@ def close(self) -> None:
             self._close()
         self._close = None
 
-    def isnull(
-        self: T_DataWithCoords, keep_attrs: bool | None = None
-    ) -> T_DataWithCoords:
+    def isnull(self, keep_attrs: bool | None = None) -> Self:
         """Test each value in the array for whether it is a missing value.
 
         Parameters
@@ -1253,9 +1248,7 @@ def isnull(
             keep_attrs=keep_attrs,
         )
 
-    def notnull(
-        self: T_DataWithCoords, keep_attrs: bool | None = None
-    ) -> T_DataWithCoords:
+    def notnull(self, keep_attrs: bool | None = None) -> Self:
         """Test each value in the array for whether it is not a missing value.
 
         Parameters
@@ -1298,7 +1291,7 @@ def notnull(
             keep_attrs=keep_attrs,
         )
 
-    def isin(self: T_DataWithCoords, test_elements: Any) -> T_DataWithCoords:
+    def isin(self, test_elements: Any) -> Self:
         """Tests each value in the array for whether it is in test elements.
 
         Parameters
@@ -1347,15 +1340,15 @@ def isin(self: T_DataWithCoords, test_elements: Any) -> T_DataWithCoords:
         )
 
     def astype(
-        self: T_DataWithCoords,
+        self,
         dtype,
         *,
         order=None,
         casting=None,
         subok=None,
         copy=None,
         keep_attrs=True,
-    ) -> T_DataWithCoords:
+    ) -> Self:
         """
         Copy of the xarray object, with data cast to a specified type.
         Leaves coordinate dtype unchanged.
@@ -1422,7 +1415,7 @@ def astype(
             dask="allowed",
         )
 
-    def __enter__(self: T_DataWithCoords) -> T_DataWithCoords:
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:

xarray/core/concat.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Hashable, Iterable
-from typing import TYPE_CHECKING, Any, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Union, overload
 
 import numpy as np
 import pandas as pd
@@ -504,8 +504,7 @@ def _dataset_concat(
 
     # case where concat dimension is a coordinate or data_var but not a dimension
     if (dim in coord_names or dim in data_names) and dim not in dim_names:
-        # TODO: Overriding type because .expand_dims has incorrect typing:
-        datasets = [cast(T_Dataset, ds.expand_dims(dim)) for ds in datasets]
+        datasets = [ds.expand_dims(dim) for ds in datasets]
 
     # determine which variables to concatenate
     concat_over, equals, concat_dim_lengths = _calc_concat_over(
@@ -708,8 +707,7 @@ def _dataarray_concat(
             if compat == "identical":
                 raise ValueError("array names not identical")
             else:
-                # TODO: Overriding type because .rename has incorrect typing:
-                arr = cast(T_DataArray, arr.rename(name))
+                arr = arr.rename(name)
         datasets.append(arr._to_temp_dataset())
 
     ds = _dataset_concat(

xarray/core/coordinates.py

@@ -23,7 +23,7 @@
     create_default_index_implicit,
 )
 from xarray.core.merge import merge_coordinates_without_align, merge_coords
-from xarray.core.types import Self, T_DataArray
+from xarray.core.types import Self, T_DataArray, T_Xarray
 from xarray.core.utils import (
     Frozen,
     ReprObject,
@@ -425,7 +425,7 @@ def __delitem__(self, key: Hashable) -> None:
         # redirect to DatasetCoordinates.__delitem__
         del self._data.coords[key]
 
-    def equals(self, other: Coordinates) -> bool:
+    def equals(self, other: Self) -> bool:
         """Two Coordinates objects are equal if they have matching variables,
         all of which are equal.
 
@@ -437,7 +437,7 @@ def equals(self, other: Coordinates) -> bool:
             return False
         return self.to_dataset().equals(other.to_dataset())
 
-    def identical(self, other: Coordinates) -> bool:
+    def identical(self, other: Self) -> bool:
         """Like equals, but also checks all variable attributes.
 
         See Also
@@ -565,9 +565,7 @@ def update(self, other: Mapping[Any, Any]) -> None:
 
         self._update_coords(coords, indexes)
 
-    def assign(
-        self, coords: Mapping | None = None, **coords_kwargs: Any
-    ) -> Coordinates:
+    def assign(self, coords: Mapping | None = None, **coords_kwargs: Any) -> Self:
         """Assign new coordinates (and indexes) to a Coordinates object, returning
         a new object with all the original coordinates in addition to the new ones.
 
@@ -656,16 +654,24 @@ def copy(
         self,
         deep: bool = False,
         memo: dict[int, Any] | None = None,
-    ) -> Coordinates:
+    ) -> Self:
         """Return a copy of this Coordinates object."""
         # do not copy indexes (may corrupt multi-coordinate indexes)
         # TODO: disable variables deepcopy? it may also be problematic when they
         # encapsulate index objects like pd.Index
         variables = {
             k: v._copy(deep=deep, memo=memo) for k, v in self.variables.items()
         }
-        return Coordinates._construct_direct(
-            coords=variables, indexes=dict(self.xindexes), dims=dict(self.sizes)
+
+        # TODO: getting an error with `self._construct_direct`, possibly because of how
+        # a subclass implements `_construct_direct`. (This was originally the same
+        # runtime code, but we switched the type definitions in #8216, which
+        # necessitates the cast.)
+        return cast(
+            Self,
+            Coordinates._construct_direct(
+                coords=variables, indexes=dict(self.xindexes), dims=dict(self.sizes)
+            ),
         )
 
 
@@ -915,9 +921,7 @@ def drop_indexed_coords(
     return Coordinates._construct_direct(coords=new_variables, indexes=new_indexes)
 
 
-def assert_coordinate_consistent(
-    obj: T_DataArray | Dataset, coords: Mapping[Any, Variable]
-) -> None:
+def assert_coordinate_consistent(obj: T_Xarray, coords: Mapping[Any, Variable]) -> None:
     """Make sure the dimension coordinate of obj is consistent with coords.
 
     obj: DataArray or Dataset