typing for numpy 1.20 (#4878)

* typing for numpy 1.20

* [skip-ci] add whats-new.rst

* update formatting

* -> np.dtype

* fix bug, use Mapping, check for dict-like

* enable typing CI

* fixes

* remove some unnecessary ignores again

Author: mathause

.github/workflows/ci-additional.yaml

@@ -161,7 +161,7 @@ jobs:
     name: Type checking (mypy)
     runs-on: "ubuntu-latest"
     needs: detect-ci-trigger
-    if: false && needs.detect-ci-trigger.outputs.triggered == 'false'
+    if: needs.detect-ci-trigger.outputs.triggered == 'false'
     defaults:
       run:
         shell: bash -l {0}

doc/whats-new.rst

@@ -190,6 +190,7 @@ Internal Changes
   in ipython (:issue:`4741`, :pull:`4742`). By `Richard Kleijn <https://github.com/rhkleijn>`_.
 - Added the ``set_close`` method to ``Dataset`` and ``DataArray`` for beckends to specify how to voluntary release
   all resources. (:pull:`#4809`), By `Alessandro Amici <https://github.com/alexamici>`_.
+- Update type hints to work with numpy v1.20 (:pull:`4878`). By `Mathias Hauser <https://github.com/mathause>`_.
 - Ensure warnings cannot be turned into exceptions in :py:func:`testing.assert_equal` and
   the other ``assert_*`` functions (:pull:`4864`). By `Mathias Hauser <https://github.com/mathause>`_.
 - Performance improvement when constructing DataArrays. Significantly speeds up repr for Datasets with large number of variables.

xarray/core/accessor_dt.py

@@ -9,6 +9,7 @@
     is_np_datetime_like,
     is_np_timedelta_like,
 )
+from .npcompat import DTypeLike
 from .pycompat import is_duck_dask_array
 
 
@@ -178,8 +179,9 @@ class Properties:
     def __init__(self, obj):
         self._obj = obj
 
-    def _tslib_field_accessor(  # type: ignore
-        name: str, docstring: str = None, dtype: np.dtype = None
+    @staticmethod
+    def _tslib_field_accessor(
+        name: str, docstring: str = None, dtype: DTypeLike = None
     ):
         def f(self, dtype=dtype):
             if dtype is None:

xarray/core/common.py

@@ -16,6 +16,7 @@
     Tuple,
     TypeVar,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -35,6 +36,8 @@
 
 if TYPE_CHECKING:
     from .dataarray import DataArray
+    from .dataset import Dataset
+    from .variable import Variable
     from .weighted import Weighted
 
 T_DataWithCoords = TypeVar("T_DataWithCoords", bound="DataWithCoords")
@@ -1501,7 +1504,26 @@ def __getitem__(self, value):
         raise NotImplementedError()
 
 
-def full_like(other, fill_value, dtype: DTypeLike = None):
+@overload
+def full_like(
+    other: "Dataset",
+    fill_value,
+    dtype: Union[DTypeLike, Mapping[Hashable, DTypeLike]] = None,
+) -> "Dataset":
+    ...
+
+
+@overload
+def full_like(other: "DataArray", fill_value, dtype: DTypeLike = None) -> "DataArray":
+    ...
+
+
+@overload
+def full_like(other: "Variable", fill_value, dtype: DTypeLike = None) -> "Variable":
+    ...
+
+
+def full_like(other, fill_value, dtype=None):
     """Return a new object with the same shape and type as a given object.
 
     Parameters
@@ -1618,15 +1640,22 @@ def full_like(other, fill_value, dtype: DTypeLike = None):
             f"fill_value must be scalar or, for datasets, a dict-like. Received {fill_value} instead."
         )
 
+    if not isinstance(other, Dataset) and isinstance(dtype, Mapping):
+        raise ValueError(
+            "'dtype' cannot be dict-like when passing a DataArray or Variable"
+        )
+
     if isinstance(other, Dataset):
         if not isinstance(fill_value, dict):
             fill_value = {k: fill_value for k in other.data_vars.keys()}
 
-        if not isinstance(dtype, dict):
-            dtype = {k: dtype for k in other.data_vars.keys()}
+        if not isinstance(dtype, Mapping):
+            dtype_ = {k: dtype for k in other.data_vars.keys()}
+        else:
+            dtype_ = dtype
 
         data_vars = {
-            k: _full_like_variable(v, fill_value.get(k, dtypes.NA), dtype.get(k, None))
+            k: _full_like_variable(v, fill_value.get(k, dtypes.NA), dtype_.get(k, None))
             for k, v in other.data_vars.items()
         }
         return Dataset(data_vars, coords=other.coords, attrs=other.attrs)

xarray/core/dataset.py

@@ -4311,7 +4311,7 @@ def dropna(
             subset = iter(self.data_vars)
 
         count = np.zeros(self.dims[dim], dtype=np.int64)
-        size = 0
+        size = np.int_(0)  # for type checking
 
         for k in subset:
             array = self._variables[k]
@@ -6370,7 +6370,7 @@ def polyfit(
         lhs = np.vander(x, order)
 
         if rcond is None:
-            rcond = x.shape[0] * np.core.finfo(x.dtype).eps
+            rcond = x.shape[0] * np.core.finfo(x.dtype).eps  # type: ignore
 
         # Weights:
         if w is not None:
@@ -6414,7 +6414,7 @@ def polyfit(
                 # deficient ranks nor does it output the "full" info (issue dask/dask#6516)
                 skipna_da = True
             elif skipna is None:
-                skipna_da = np.any(da.isnull())
+                skipna_da = bool(np.any(da.isnull()))
 
             dims_to_stack = [dimname for dimname in da.dims if dimname != dim]
             stacked_coords: Dict[Hashable, DataArray] = {}

xarray/core/formatting.py

@@ -189,9 +189,8 @@ def format_array_flat(array, max_width: int):
         (max_possibly_relevant < array.size) or (cum_len > max_width).any()
     ):
         padding = " ... "
-        count = min(
-            array.size, max(np.argmax(cum_len + len(padding) - 1 > max_width), 2)
-        )
+        max_len = max(np.argmax(cum_len + len(padding) - 1 > max_width), 2)  # type: ignore
+        count = min(array.size, max_len)
     else:
         count = array.size
         padding = "" if (count <= 1) else " "

xarray/core/indexing.py

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from contextlib import suppress
 from datetime import timedelta
-from typing import Any, Callable, Iterable, Sequence, Tuple, Union
+from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -1010,7 +1010,7 @@ def _decompose_outer_indexer(
         return indexer, BasicIndexer(())
     assert isinstance(indexer, (OuterIndexer, BasicIndexer))
 
-    backend_indexer = []
+    backend_indexer: List[Any] = []
     np_indexer = []
     # make indexer positive
     pos_indexer = []
@@ -1397,17 +1397,17 @@ def __init__(self, array: Any, dtype: DTypeLike = None):
         self.array = utils.safe_cast_to_index(array)
         if dtype is None:
             if isinstance(array, pd.PeriodIndex):
-                dtype = np.dtype("O")
+                dtype_ = np.dtype("O")
             elif hasattr(array, "categories"):
                 # category isn't a real numpy dtype
-                dtype = array.categories.dtype
+                dtype_ = array.categories.dtype
             elif not utils.is_valid_numpy_dtype(array.dtype):
-                dtype = np.dtype("O")
+                dtype_ = np.dtype("O")
             else:
-                dtype = array.dtype
+                dtype_ = array.dtype
         else:
-            dtype = np.dtype(dtype)
-        self._dtype = dtype
+            dtype_ = np.dtype(dtype)
+        self._dtype = dtype_
 
     @property
     def dtype(self) -> np.dtype:

xarray/core/npcompat.py

@@ -75,12 +75,12 @@ def moveaxis(a, source, destination):
     return result
 
 
-# Type annotations stubs.
+# Type annotations stubs
 try:
     from numpy.typing import DTypeLike
 except ImportError:
     # fall back for numpy < 1.20
-    DTypeLike = Union[np.dtype, str]
+    DTypeLike = Union[np.dtype, str]  # type: ignore
 
 
 # from dask/array/utils.py

xarray/core/nputils.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import pandas as pd
-from numpy.core.multiarray import normalize_axis_index
+from numpy.core.multiarray import normalize_axis_index  # type: ignore
 
 try:
     import bottleneck as bn

xarray/tests/test_cftime_offsets.py

@@ -479,7 +479,7 @@ def test_minus_offset(a, b):
 
 @pytest.mark.parametrize(
     ("a", "b"),
-    list(zip(np.roll(_EQ_TESTS_A, 1), _EQ_TESTS_B))
+    list(zip(np.roll(_EQ_TESTS_A, 1), _EQ_TESTS_B))  # type: ignore
     + [(YearEnd(month=1), YearEnd(month=2))],
     ids=_id_func,
 )