(fix): more typing

Author: ilan-gold

properties/test_pandas_roundtrip.py

@@ -3,12 +3,14 @@
 """
 
 from functools import partial
+from typing import cast
 
 import numpy as np
 import pandas as pd
 import pytest
 
 import xarray as xr
+from xarray.core.dataset import Dataset
 
 pytest.importorskip("hypothesis")
 import hypothesis.extra.numpy as npst  # isort:skip
@@ -88,7 +90,7 @@ def test_roundtrip_dataarray(data, arr) -> None:
 
 
 @given(datasets_1d_vars())
-def test_roundtrip_dataset(dataset) -> None:
+def test_roundtrip_dataset(dataset: Dataset) -> None:
     df = dataset.to_dataframe()
     assert isinstance(df, pd.DataFrame)
     roundtripped = xr.Dataset(df)
@@ -119,7 +121,7 @@ def test_roundtrip_pandas_dataframe(df) -> None:
     df.columns.name = "cols"
     arr = xr.DataArray(df)
     roundtripped = arr.to_pandas()
-    pd.testing.assert_frame_equal(df, roundtripped)
+    pd.testing.assert_frame_equal(df, cast(pd.DataFrame, roundtripped))
     xr.testing.assert_identical(arr, roundtripped.to_xarray())
 
 
@@ -143,8 +145,8 @@ def test_roundtrip_pandas_dataframe_datetime(df) -> None:
         pd.arrays.IntervalArray(
             [pd.Interval(0, 1), pd.Interval(1, 5), pd.Interval(2, 6)]
         ),
-        pd.arrays.TimedeltaArray._from_sequence(pd.TimedeltaIndex(["1h", "2h", "3h"])),
-        pd.arrays.DatetimeArray._from_sequence(
+        pd.arrays.TimedeltaArray._from_sequence(pd.TimedeltaIndex(["1h", "2h", "3h"])),  # type: ignore[attr-defined]
+        pd.arrays.DatetimeArray._from_sequence(  # type: ignore[attr-defined]
             pd.DatetimeIndex(["2023-01-01", "2023-01-02", "2023-01-03"], freq="D")
         ),
         np.array([1, 2, 3], dtype="int64"),

xarray/core/dtypes.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import functools
-from typing import TYPE_CHECKING
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 from pandas.api.types import is_extension_array_dtype
@@ -53,7 +54,9 @@ def __eq__(self, other):
 )
 
 
-def maybe_promote(dtype: np.dtype | ExtensionDtype) -> tuple[np.dtype, Any]:
+def maybe_promote(
+    dtype: np.dtype | ExtensionDtype,
+) -> tuple[np.dtype | ExtensionDtype, Any]:
     """Simpler equivalent of pandas.core.common._maybe_promote
 
     Parameters
@@ -70,7 +73,9 @@ def maybe_promote(dtype: np.dtype | ExtensionDtype) -> tuple[np.dtype, Any]:
     fill_value: Any
     if is_extension_array_dtype(dtype):
         return dtype, dtype.na_value
-    elif HAS_STRING_DTYPE and np.issubdtype(dtype, np.dtypes.StringDType()):
+    else:
+        dtype = cast(np.dtype, dtype)
+    if 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
@@ -251,7 +256,7 @@ def maybe_promote_to_variable_width(
 
 
 def should_promote_to_object(
-    arrays_and_dtypes: np.typing.ArrayLike | np.typing.DTypeLike, xp
+    arrays_and_dtypes: Iterable[np.typing.ArrayLike | np.typing.DTypeLike], xp
 ) -> bool:
     """
     Test whether the given arrays_and_dtypes, when evaluated individually, match the
@@ -281,7 +286,9 @@ def should_promote_to_object(
 
 
 def result_type(
-    *arrays_and_dtypes: np.typing.ArrayLike | np.typing.DTypeLike,
+    *arrays_and_dtypes: list[
+        np.typing.ArrayLike | np.typing.DTypeLike | ExtensionDtype
+    ],
     xp=None,
 ) -> np.dtype:
     """Like np.result_type, but with type promotion rules matching pandas.
@@ -314,7 +321,7 @@ def result_type(
                 maybe_promote_to_variable_width,
                 # let extension arrays handle their own str/bytes
                 should_return_str_or_bytes=any(
-                    map(is_extension_array_dtype, arrays_and_dtypes)
+                    map(is_extension_array_dtype, arrays_and_dtypes)  # type: ignore[arg-type]
                 ),
             ),
             arrays_and_dtypes,

xarray/core/extension_array.py

@@ -58,7 +58,7 @@ def __extension_duck_array__issubdtype(
 
 @implements("astype")  # np.astype was added in 2.1.0, but we only require >=1.24
 def __extension_duck_array__astype(
-    array_or_scalar: np.typing.ArrayLike,
+    array_or_scalar: T_ExtensionArray,
     dtype: DTypeLikeSave,
     order: str = "K",
     casting: str = "unsafe",
@@ -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)
+            is_extension_array_dtype(array_or_scalar) or is_extension_array_dtype(dtype)  # type: ignore[arg-dtype]
         )
         or casting != "unsafe"
         or not subok
@@ -81,7 +81,8 @@ def __extension_duck_array__astype(
 
 @implements(np.asarray)
 def __extension_duck_array__asarray(
-    array_or_scalar: np.typing.ArrayLike, dtype: DTypeLikeSave = None
+    array_or_scalar: np.typing.ArrayLike | T_ExtensionArray,
+    dtype: DTypeLikeSave | None = None,
 ) -> ExtensionArray:
     if not is_extension_array_dtype(dtype):
         return NotImplemented
@@ -90,7 +91,9 @@ def __extension_duck_array__asarray(
 
 
 def as_extension_array(
-    array_or_scalar: np.typing.ArrayLike, dtype: ExtensionDtype, copy: bool = False
+    array_or_scalar: np.typing.ArrayLike | T_ExtensionArray,
+    dtype: ExtensionDtype | DTypeLikeSave | None,
+    copy: bool = False,
 ) -> ExtensionArray:
     if is_scalar(array_or_scalar):
         return dtype.construct_array_type()._from_sequence(  # type: ignore[attr-defined]