Skip to content

Commit

Permalink
typing for numpy 1.20 (#4878)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
mathause authored Feb 23, 2021
1 parent 348eb48 commit 0f65307
Show file tree
Hide file tree
Showing 12 changed files with 62 additions and 25 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/ci-additional.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
1 change: 1 addition & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
6 changes: 4 additions & 2 deletions xarray/core/accessor_dt.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
is_np_datetime_like,
is_np_timedelta_like,
)
from .npcompat import DTypeLike
from .pycompat import is_duck_dask_array


Expand Down Expand Up @@ -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:
Expand Down
37 changes: 33 additions & 4 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
Tuple,
TypeVar,
Union,
overload,
)

import numpy as np
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
6 changes: 3 additions & 3 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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] = {}
Expand Down
5 changes: 2 additions & 3 deletions xarray/core/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 " "
Expand Down
16 changes: 8 additions & 8 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 = []
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions xarray/core/npcompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion xarray/core/nputils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion xarray/tests/test_cftime_offsets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
3 changes: 3 additions & 0 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -4200,6 +4200,9 @@ def test_full_like(self):
assert expect.dtype == bool
assert_identical(expect, actual)

with pytest.raises(ValueError, match="'dtype' cannot be dict-like"):
full_like(da, fill_value=True, dtype={"x": bool})

def test_dot(self):
x = np.linspace(-3, 3, 6)
y = np.linspace(-3, 3, 5)
Expand Down
3 changes: 3 additions & 0 deletions xarray/tests/test_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -2248,6 +2248,9 @@ def test_full_like(self):
with raises_regex(ValueError, "must be scalar"):
full_like(orig, [1.0, 2.0])

with pytest.raises(ValueError, match="'dtype' cannot be dict-like"):
full_like(orig, True, dtype={"x": bool})

@requires_dask
def test_full_like_dask(self):
orig = Variable(
Expand Down

0 comments on commit 0f65307

Please sign in to comment.