{full,zeros,ones}_like typing

[headtr1ck]

(partial) typing for functions full_like, zeros_like, ones_like.

I could not figure out how to properly use TypeVars so many things are "hardcoded" with overloads.

I have added a DTypeLikeSave to npcompat, not sure that this file is supposed to be edited.

Problem1: TypeVar["T", Dataset, DataArray, Variable] can only be one of these three, but never with Union[Dataset, DataArray] which is used in several other places in xarray.
Problem2: The official mypy support says: use TypeVar["T", bound=Union[Dataset, DataArray, Variable] but the the isinstance(obj, Dataset) could not be correctly resolved (is that a mypy issue?).

So if anyone can get it to work with TypeVars, feel free to change it. :)

[max-sixty]

Looks great, thanks a lot @headtr1ck .

I couldn't seem to push to the PR, but I added some annotations to our tests so we could also test the type annotations:

commit cb4c5dd2dac49230220beda3fa72b286c8849ab4
Author: Maximilian Roos <[email protected]>
Date:   Sun May 15 14:41:46 2022 -0700

    Type check more tests

diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 263237d9..950f15e9 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -5559,7 +5559,7 @@ def test_binary_op_join_setting(self):
             actual = ds1 + ds2
             assert_equal(actual, expected)
 
-    def test_full_like(self):
+    def test_full_like(self) -> None:
         # For more thorough tests, see test_variable.py
         # Note: testing data_vars with mismatched dtypes
         ds = Dataset(
@@ -5572,8 +5572,9 @@ def test_full_like(self):
         actual = full_like(ds, 2)
 
         expected = ds.copy(deep=True)
-        expected["d1"].values = [2, 2, 2]
-        expected["d2"].values = [2.0, 2.0, 2.0]
+        # https://github.com/python/mypy/issues/3004
+        expected["d1"].values = [2, 2, 2]  # type: ignore
+        expected["d2"].values = [2.0, 2.0, 2.0]  # type: ignore
         assert expected["d1"].dtype == int
         assert expected["d2"].dtype == float
         assert_identical(expected, actual)
@@ -5581,8 +5582,8 @@ def test_full_like(self):
         # override dtype
         actual = full_like(ds, fill_value=True, dtype=bool)
         expected = ds.copy(deep=True)
-        expected["d1"].values = [True, True, True]
-        expected["d2"].values = [True, True, True]
+        expected["d1"].values = [True, True, True]  # type: ignore
+        expected["d2"].values = [True, True, True]  # type: ignore
         assert expected["d1"].dtype == bool
         assert expected["d2"].dtype == bool
         assert_identical(expected, actual)
@@ -5788,7 +5789,7 @@ def test_ipython_key_completion(self):
             ds.data_vars[item]  # should not raise
         assert sorted(actual) == sorted(expected)
 
-    def test_polyfit_output(self):
+    def test_polyfit_output(self) -> None:
         ds = create_test_data(seed=1)
 
         out = ds.polyfit("dim2", 2, full=False)
@@ -5801,7 +5802,7 @@ def test_polyfit_output(self):
         out = ds.polyfit("time", 2)
         assert len(out.data_vars) == 0
 
-    def test_polyfit_warnings(self):
+    def test_polyfit_warnings(self) -> None:
         ds = create_test_data(seed=1)
 
         with warnings.catch_warnings(record=True) as ws:
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 0168f19b..886b0360 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -2480,7 +2480,7 @@ def test_datetime(self):
         assert np.ndarray == type(actual)
         assert np.dtype("datetime64[ns]") == actual.dtype
 
-    def test_full_like(self):
+    def test_full_like(self) -> None:
         # For more thorough tests, see test_variable.py
         orig = Variable(
             dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"}
@@ -2503,7 +2503,7 @@ def test_full_like(self):
             full_like(orig, True, dtype={"x": bool})
 
     @requires_dask
-    def test_full_like_dask(self):
+    def test_full_like_dask(self) -> None:
         orig = Variable(
             dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"}
         ).chunk(((1, 1), (2,)))
@@ -2534,14 +2534,14 @@ def check(actual, expect_dtype, expect_values):
             else:
                 assert not isinstance(v, np.ndarray)
 
-    def test_zeros_like(self):
+    def test_zeros_like(self) -> None:
         orig = Variable(
             dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"}
         )
         assert_identical(zeros_like(orig), full_like(orig, 0))
         assert_identical(zeros_like(orig, dtype=int), full_like(orig, 0, dtype=int))
 
-    def test_ones_like(self):
+    def test_ones_like(self) -> None:
         orig = Variable(
             dims=("x", "y"), data=[[1.5, 2.0], [3.1, 4.3]], attrs={"foo": "bar"}
         )

[max-sixty]

For these ones which are generic over DataArray and Variable, we could define a TypeVar for them (but no need in this PR, and I'm not even sure whether it's a broader issue than just the *_like.

[max-sixty]

Do you know why this overload is required? I had thought that DataArray couldn't take a DTypeMaybeMapping? (and same principle for the next two)

[max-sixty]

I tried to adjust these locally. An issue I hit is that so much of our typing is on Dataset rather than T_Dataset and changing one seems to require changing its dependencies, which makes it difficult to change gradually. 20 minutes into me trying to change this, I had 31 errors in 6 files!

To the extent you know which parts are Perfect vs Not-perfect-but-required-to-pass: If you want to add a comment for the ones the latter, that will make it easier for future travelers to know why things are as they are and hopefully change them.

[headtr1ck]

I was struggling a lot with TypeVar definitions.
Some other functions (like polyval) are calling zeros_like with DataArray | Dataset.

This means that a "simple" TypeVar["T", DataArray, Dataset] will not work.

I tried using TypeVar["T", bound=DataArray|Dataset] (recommended by some mypy people) but then the if isinstance(x Dataset) were causing problems (still not sure if that is a mypy bug or intended, my TypeVar knowledge is not good enough for that)...

So this was the only solution I could get to work.

[headtr1ck]

Additionally, the code will actually create a plain e.g. DataArray, so typevars with bounds are actually wrong here.

[max-sixty]

Yeah. I think the proposed code is a big upgrade, and we can refine towards perfection in the future...

[max-sixty]

FWIW I found this SO answer helpful in clarifying the difference — i saw that I had upvoted it before — but I'm still not confident in how we should design these methods.

[max-sixty]

This is the kind of func that would be nice at some point to make generic; with the proposed code we lose whether it's a Dataset vs. DataArray. (fine to add as a comment / TODO tho)

[headtr1ck]

I failed to make it work with Typevars since it is called with DataArray|Dataset :(

[max-sixty]

Thanks a lot @headtr1ck ! I will hit the auto-merge button!

[max-sixty]

One thought for any future travelers — if we want to move from Union[DataArray, Dataset] to generic over DataArray & Dataset without changing everything at once (as was the issue here) we can try adding the generic type that methods accept without removing existing types, e.g. Union[DataArray, Dataset, T_Xarray].