(fix): more typing!

Author: ilan-gold

properties/test_pandas_roundtrip.py

@@ -93,7 +93,7 @@ def test_roundtrip_dataarray(data, arr) -> None:
 def test_roundtrip_dataset(dataset: Dataset) -> None:
     df = dataset.to_dataframe()
     assert isinstance(df, pd.DataFrame)
-    roundtripped = xr.Dataset(df)
+    roundtripped = xr.Dataset.from_dataframe(df)
     xr.testing.assert_identical(dataset, roundtripped)
 
 
@@ -103,7 +103,7 @@ def test_roundtrip_pandas_series(ser, ix_name) -> None:
     ser.index.name = ix_name
     arr = xr.DataArray(ser)
     roundtripped = arr.to_pandas()
-    pd.testing.assert_series_equal(ser, roundtripped)
+    pd.testing.assert_series_equal(ser, roundtripped)  # type: ignore[arg-type]
     xr.testing.assert_identical(arr, roundtripped.to_xarray())
 
 

xarray/core/dtypes.py

@@ -2,9 +2,10 @@
 
 import functools
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING, TypeVar, cast
 
 import numpy as np
+from pandas.api.extensions import ExtensionDtype
 from pandas.api.types import is_extension_array_dtype
 
 from xarray.compat import array_api_compat, npcompat
@@ -14,7 +15,6 @@
 if TYPE_CHECKING:
     from typing import Any
 
-    from pandas.api.extensions import ExtensionDtype
 
 # Use as a sentinel value to indicate a dtype appropriate NA value.
 NA = utils.ReprObject("<NA>")
@@ -53,10 +53,10 @@ def __eq__(self, other):
     (np.bytes_, np.str_),  # numpy promotes to unicode
 )
 
+T_dtype = TypeVar("T_dtype", np.dtype, ExtensionDtype)
 
-def maybe_promote(
-    dtype: np.dtype | ExtensionDtype,
-) -> tuple[np.dtype | ExtensionDtype, Any]:
+
+def maybe_promote(dtype: T_dtype) -> tuple[T_dtype, Any]:
     """Simpler equivalent of pandas.core.common._maybe_promote
 
     Parameters
@@ -72,10 +72,12 @@ def maybe_promote(
     dtype_: np.typing.DTypeLike
     fill_value: Any
     if is_extension_array_dtype(dtype):
-        return dtype, dtype.na_value
-    else:
-        dtype = cast(np.dtype, dtype)
-    if HAS_STRING_DTYPE and np.issubdtype(dtype, np.dtypes.StringDType()):
+        return dtype, cast(ExtensionDtype, dtype).na_value  # type: ignore[redundant-cast]
+    if not isinstance(dtype, np.dtype):
+        raise TypeError(
+            f"dtype {dtype} must be one of an extension array dtype or numpy dtype"
+        )
+    elif HAS_STRING_DTYPE and np.issubdtype(dtype, np.dtypes.StringDType()):
         # for now, we always promote string dtypes to object for consistency with existing behavior
         # TODO: refactor this once we have a better way to handle numpy vlen-string dtypes
         dtype_ = object
@@ -235,10 +237,14 @@ def isdtype(dtype, kind: str | tuple[str, ...], xp=None) -> bool:
 
 
 def maybe_promote_to_variable_width(
-    array_or_dtype: np.typing.ArrayLike | np.typing.DTypeLike,
+    array_or_dtype: np.typing.ArrayLike
+    | np.typing.DTypeLike
+    | ExtensionDtype
+    | str
+    | bytes,
     *,
     should_return_str_or_bytes: bool = False,
-) -> np.typing.ArrayLike | np.typing.DTypeLike:
+) -> np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype:
     if isinstance(array_or_dtype, str | bytes):
         if should_return_str_or_bytes:
             return array_or_dtype
@@ -256,7 +262,10 @@ def maybe_promote_to_variable_width(
 
 
 def should_promote_to_object(
-    arrays_and_dtypes: Iterable[np.typing.ArrayLike | np.typing.DTypeLike], xp
+    arrays_and_dtypes: Iterable[
+        np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype
+    ],
+    xp,
 ) -> bool:
     """
     Test whether the given arrays_and_dtypes, when evaluated individually, match the
@@ -286,9 +295,7 @@ def should_promote_to_object(
 
 
 def result_type(
-    *arrays_and_dtypes: list[
-        np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype
-    ],
+    *arrays_and_dtypes: np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype,
     xp=None,
 ) -> np.dtype:
     """Like np.result_type, but with type promotion rules matching pandas.

xarray/core/extension_array.py

@@ -68,7 +68,7 @@ def __extension_duck_array__astype(
 ) -> ExtensionArray:
     if (
         not (
-            is_extension_array_dtype(array_or_scalar) or is_extension_array_dtype(dtype)  # type: ignore[arg-dtype]
+            is_extension_array_dtype(array_or_scalar) or is_extension_array_dtype(dtype)
         )
         or casting != "unsafe"
         or not subok
@@ -96,21 +96,23 @@ def as_extension_array(
     copy: bool = False,
 ) -> ExtensionArray:
     if is_scalar(array_or_scalar):
-        return dtype.construct_array_type()._from_sequence(  # type: ignore[attr-defined]
+        return dtype.construct_array_type()._from_sequence(  # type: ignore[union-attr]
             [array_or_scalar], dtype=dtype
         )
     else:
-        return array_or_scalar.astype(dtype, copy=copy)
+        return array_or_scalar.astype(dtype, copy=copy)  # type: ignore[union-attr]
 
 
 @implements(np.result_type)
 def __extension_duck_array__result_type(
-    *arrays_and_dtypes: np.typing.ArrayLike | np.typing.DTypeLike,
+    *arrays_and_dtypes: list[
+        np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype | ExtensionArray
+    ],
 ) -> DtypeObj:
     extension_arrays_and_dtypes: list[ExtensionDtype | ExtensionArray] = [
-        x
+        cast(ExtensionDtype | ExtensionArray, x)
         for x in arrays_and_dtypes
-        if is_extension_array_dtype(x)  # type: ignore[arg-type, misc]
+        if is_extension_array_dtype(x)
     ]
     if not extension_arrays_and_dtypes:
         return NotImplemented
@@ -120,15 +122,17 @@ def __extension_duck_array__result_type(
         for x in extension_arrays_and_dtypes
     ]
     scalars: list[Scalar] = [
-        x for x in arrays_and_dtypes if is_scalar(x) and x not in {pd.NA, np.nan}
+        cast(Scalar, x)
+        for x in arrays_and_dtypes
+        if is_scalar(x) and x not in {pd.NA, np.nan}
     ]
     # other_stuff could include:
     # - arrays such as pd.ABCSeries, np.ndarray, or other array-api duck arrays
     # - dtypes such as pd.DtypeObj, np.dtype, or other array-api duck dtypes
     other_stuff = [
         x
         for x in arrays_and_dtypes
-        if not is_extension_array_dtype(x) and not is_scalar(x)  # type: ignore[arg-type, misc]
+        if not is_extension_array_dtype(x) and not is_scalar(x)
     ]
     # We implement one special case: when possible, preserve Categoricals (avoid promoting
     # to object) by merging the categories of all given Categoricals + scalars + NA.
@@ -178,14 +182,14 @@ def __extension_duck_array__concatenate(
 
 @implements(np.where)
 def __extension_duck_array__where(
-    condition: T_ExtensionArray | np.ArrayLike,
+    condition: T_ExtensionArray | np.typing.ArrayLike,
     x: T_ExtensionArray,
-    y: T_ExtensionArray | np.ArrayLike,
+    y: T_ExtensionArray | np.typing.ArrayLike,
 ) -> T_ExtensionArray:
     return cast(T_ExtensionArray, pd.Series(x).where(condition, y).array)  # type: ignore[arg-type]
 
 
-def _replace_duck(args, replacer: Callable[[PandasExtensionArray]]) -> list:
+def _replace_duck(args, replacer: Callable[[PandasExtensionArray], list]) -> list:
     args_as_list = list(args)
     for index, value in enumerate(args_as_list):
         if isinstance(value, PandasExtensionArray):