From ae81de3ba97fede5b384088293e1f73c67672f81 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 15 Oct 2024 20:55:43 -0400 Subject: [PATCH 1/5] type hints for datatree ops tests --- xarray/tests/test_datatree.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 97fe8212c86..eca990854d2 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1769,7 +1769,7 @@ def test_subtree(self): class TestOps: - def test_unary_op(self): + def test_unary_op(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataTree.from_dict({"/": ds1, "/subnode": ds2}) @@ -1779,7 +1779,7 @@ def test_unary_op(self): result = -dt assert_equal(result, expected) - def test_unary_op_inherited_coords(self): + def test_unary_op_inherited_coords(self) -> None: tree = DataTree(xr.Dataset(coords={"x": [1, 2, 3]})) tree["/foo"] = DataTree(xr.Dataset({"bar": ("x", [4, 5, 6])})) actual = -tree @@ -1792,7 +1792,7 @@ def test_unary_op_inherited_coords(self): expected["/foo/bar"].data = np.array([-4, -5, -6]) assert_identical(actual, expected) - def test_binary_op_on_int(self): + def test_binary_op_on_int(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataTree.from_dict({"/": ds1, "/subnode": ds2}) @@ -1802,7 +1802,7 @@ def test_binary_op_on_int(self): result = dt * 5 assert_equal(result, expected) - def test_binary_op_on_dataarray(self): + def test_binary_op_on_dataarray(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataTree.from_dict( @@ -1824,7 +1824,7 @@ def test_binary_op_on_dataarray(self): result = dt * other_da assert_equal(result, expected) - def test_binary_op_on_dataset(self): + def test_binary_op_on_dataset(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataTree.from_dict( @@ -1846,7 +1846,7 @@ def test_binary_op_on_dataset(self): result = dt * other_ds assert_equal(result, expected) - def test_binary_op_on_datatree(self): + def test_binary_op_on_datatree(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) @@ -1857,7 +1857,7 @@ def test_binary_op_on_datatree(self): result = dt * dt assert_equal(result, expected) - def test_arithmetic_inherited_coords(self): + def test_arithmetic_inherited_coords(self) -> None: tree = DataTree(xr.Dataset(coords={"x": [1, 2, 3]})) tree["/foo"] = DataTree(xr.Dataset({"bar": ("x", [4, 5, 6])})) actual = 2 * tree @@ -1869,7 +1869,7 @@ def test_arithmetic_inherited_coords(self): expected["/foo/bar"].data = np.array([8, 10, 12]) assert_identical(actual, expected) - def test_binary_op_commutativity_with_dataset(self): + def test_binary_op_commutativity_with_dataset(self) -> None: # regression test for #9365 ds1 = xr.Dataset({"a": [5], "b": [3]}) @@ -1893,7 +1893,7 @@ def test_binary_op_commutativity_with_dataset(self): result = other_ds * dt assert_equal(result, expected) - def test_inplace_binary_op(self): + def test_inplace_binary_op(self) -> None: ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataTree.from_dict({"/": ds1, "/subnode": ds2}) @@ -1903,7 +1903,7 @@ def test_inplace_binary_op(self): dt += 1 assert_equal(dt, expected) - def test_dont_broadcast_single_node_tree(self): + def test_dont_broadcast_single_node_tree(self) -> None: # regression test for https://github.com/pydata/xarray/issues/9365#issuecomment-2291622577 ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) From e2d1f310ccf807304038a191f3246900f43d63a3 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 15 Oct 2024 20:56:41 -0400 Subject: [PATCH 2/5] type hints for datatree aggregations tests --- xarray/tests/test_datatree.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index eca990854d2..9cb8754a960 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1690,7 +1690,7 @@ def test_sel(self): class TestAggregations: - def test_reduce_method(self): + def test_reduce_method(self) -> None: ds = xr.Dataset({"a": ("x", [False, True, False])}) dt = DataTree.from_dict({"/": ds, "/results": ds}) @@ -1699,7 +1699,7 @@ def test_reduce_method(self): result = dt.any() assert_equal(result, expected) - def test_nan_reduce_method(self): + def test_nan_reduce_method(self) -> None: ds = xr.Dataset({"a": ("x", [1, 2, 3])}) dt = DataTree.from_dict({"/": ds, "/results": ds}) @@ -1708,7 +1708,7 @@ def test_nan_reduce_method(self): result = dt.mean() assert_equal(result, expected) - def test_cum_method(self): + def test_cum_method(self) -> None: ds = xr.Dataset({"a": ("x", [1, 2, 3])}) dt = DataTree.from_dict({"/": ds, "/results": ds}) @@ -1722,7 +1722,7 @@ def test_cum_method(self): result = dt.cumsum() assert_equal(result, expected) - def test_dim_argument(self): + def test_dim_argument(self) -> None: dt = DataTree.from_dict( { "/a": xr.Dataset({"A": ("x", [1, 2])}), @@ -1757,7 +1757,7 @@ def test_dim_argument(self): ): dt.mean("invalid") - def test_subtree(self): + def test_subtree(self) -> None: tree = DataTree.from_dict( { "/child": Dataset({"a": ("x", [1, 2])}), From 8f28e579994cde0acd2bd39a638a76623b7c1a55 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 15 Oct 2024 20:57:17 -0400 Subject: [PATCH 3/5] type hints for datatree indexing tests --- xarray/tests/test_datatree.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 9cb8754a960..8db6f01134e 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1584,7 +1584,7 @@ def test_filter(self): class TestIndexing: - def test_isel_siblings(self): + def test_isel_siblings(self) -> None: tree = DataTree.from_dict( { "/first": xr.Dataset({"a": ("x", [1, 2])}), @@ -1616,7 +1616,7 @@ def test_isel_siblings(self): actual = tree.isel(x=slice(None)) assert_equal(actual, tree) - def test_isel_inherited(self): + def test_isel_inherited(self) -> None: tree = DataTree.from_dict( { "/": xr.Dataset(coords={"x": [1, 2]}), @@ -1664,7 +1664,7 @@ def test_isel_inherited(self): ) assert_identical(actual, expected) - def test_sel(self): + def test_sel(self) -> None: tree = DataTree.from_dict( { "/first": xr.Dataset({"a": ("x", [1, 2, 3])}, coords={"x": [1, 2, 3]}), From 1fc31ff83002d95b916b16b19d3dd03b6ff688aa Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 15 Oct 2024 21:00:57 -0400 Subject: [PATCH 4/5] type hint a lot more tests --- xarray/tests/test_datatree.py | 114 +++++++++++++++++----------------- 1 file changed, 57 insertions(+), 57 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 8db6f01134e..692f57cc252 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -472,7 +472,7 @@ def test_copy_with_data(self, create_test_datatree): class TestSetItem: - def test_setitem_new_child_node(self): + def test_setitem_new_child_node(self) -> None: john = DataTree(name="john") mary = DataTree(name="mary") john["mary"] = mary @@ -481,12 +481,12 @@ def test_setitem_new_child_node(self): assert grafted_mary.parent is john assert grafted_mary.name == "mary" - def test_setitem_unnamed_child_node_becomes_named(self): + def test_setitem_unnamed_child_node_becomes_named(self) -> None: john2 = DataTree(name="john2") john2["sonny"] = DataTree() assert john2["sonny"].name == "sonny" - def test_setitem_new_grandchild_node(self): + def test_setitem_new_grandchild_node(self) -> None: john = DataTree.from_dict({"/Mary/Rose": DataTree()}) new_rose = DataTree(dataset=xr.Dataset({"x": 0})) john["Mary/Rose"] = new_rose @@ -495,20 +495,20 @@ def test_setitem_new_grandchild_node(self): assert grafted_rose.parent is john["/Mary"] assert grafted_rose.name == "Rose" - def test_grafted_subtree_retains_name(self): + def test_grafted_subtree_retains_name(self) -> None: subtree = DataTree(name="original_subtree_name") root = DataTree(name="root") root["new_subtree_name"] = subtree # noqa assert subtree.name == "original_subtree_name" - def test_setitem_new_empty_node(self): + def test_setitem_new_empty_node(self) -> None: john = DataTree(name="john") john["mary"] = DataTree() mary = john["mary"] assert isinstance(mary, DataTree) assert_identical(mary.to_dataset(), xr.Dataset()) - def test_setitem_overwrite_data_in_node_with_none(self): + def test_setitem_overwrite_data_in_node_with_none(self) -> None: john = DataTree.from_dict({"/mary": xr.Dataset()}, name="john") john["mary"] = DataTree() @@ -519,49 +519,49 @@ def test_setitem_overwrite_data_in_node_with_none(self): john["."] = DataTree() @pytest.mark.xfail(reason="assigning Datasets doesn't yet create new nodes") - def test_setitem_dataset_on_this_node(self): + def test_setitem_dataset_on_this_node(self) -> None: data = xr.Dataset({"temp": [0, 50]}) results = DataTree(name="results") results["."] = data assert_identical(results.to_dataset(), data) - def test_setitem_dataset_as_new_node(self): + def test_setitem_dataset_as_new_node(self) -> None: data = xr.Dataset({"temp": [0, 50]}) folder1 = DataTree(name="folder1") folder1["results"] = data assert_identical(folder1["results"].to_dataset(), data) - def test_setitem_dataset_as_new_node_requiring_intermediate_nodes(self): + def test_setitem_dataset_as_new_node_requiring_intermediate_nodes(self) -> None: data = xr.Dataset({"temp": [0, 50]}) folder1 = DataTree(name="folder1") folder1["results/highres"] = data assert_identical(folder1["results/highres"].to_dataset(), data) - def test_setitem_named_dataarray(self): + def test_setitem_named_dataarray(self) -> None: da = xr.DataArray(name="temp", data=[0, 50]) folder1 = DataTree(name="folder1") folder1["results"] = da expected = da.rename("results") assert_equal(folder1["results"], expected) - def test_setitem_unnamed_dataarray(self): + def test_setitem_unnamed_dataarray(self) -> None: data = xr.DataArray([0, 50]) folder1 = DataTree(name="folder1") folder1["results"] = data assert_equal(folder1["results"], data) - def test_setitem_variable(self): + def test_setitem_variable(self) -> None: var = xr.Variable(data=[0, 50], dims="x") folder1 = DataTree(name="folder1") folder1["results"] = var assert_equal(folder1["results"], xr.DataArray(var)) - def test_setitem_coerce_to_dataarray(self): + def test_setitem_coerce_to_dataarray(self) -> None: folder1 = DataTree(name="folder1") folder1["results"] = 0 assert_equal(folder1["results"], xr.DataArray(0)) - def test_setitem_add_new_variable_to_empty_node(self): + def test_setitem_add_new_variable_to_empty_node(self) -> None: results = DataTree(name="results") results["pressure"] = xr.DataArray(data=[2, 3]) assert "pressure" in results.dataset @@ -575,7 +575,7 @@ def test_setitem_add_new_variable_to_empty_node(self): results_with_path["highres/temp"] = xr.Variable(data=[10, 11], dims=["x"]) assert "temp" in results_with_path["highres"].dataset - def test_setitem_dataarray_replace_existing_node(self): + def test_setitem_dataarray_replace_existing_node(self) -> None: t = xr.Dataset({"temp": [0, 50]}) results = DataTree(name="results", dataset=t) p = xr.DataArray(data=[2, 3]) @@ -585,7 +585,7 @@ def test_setitem_dataarray_replace_existing_node(self): class TestCoords: - def test_properties(self): + def test_properties(self) -> None: # use int64 for repr consistency on windows ds = Dataset( data_vars={ @@ -649,7 +649,7 @@ def test_properties(self): "b": np.dtype("int64"), } - def test_modify(self): + def test_modify(self) -> None: ds = Dataset( data_vars={ "foo": (["x", "y"], np.random.randn(2, 3)), @@ -709,7 +709,7 @@ def test_modify(self): dt2 = DataTree(dataset=dt.coords) assert_identical(dt2.coords, dt.coords) - def test_inherited(self): + def test_inherited(self) -> None: ds = Dataset( data_vars={ "foo": (["x", "y"], np.random.randn(2, 3)), @@ -750,7 +750,7 @@ def test_inherited(self): # assert_identical(expected, actual) -def test_delitem(): +def test_delitem() -> None: ds = Dataset({"a": 0}, coords={"x": ("x", [1, 2]), "z": "a"}) dt = DataTree(ds, children={"c": DataTree()}) @@ -786,7 +786,7 @@ def test_delitem(): class TestTreeFromDict: - def test_data_in_root(self): + def test_data_in_root(self) -> None: dat = xr.Dataset() dt = DataTree.from_dict({"/": dat}) assert dt.name is None @@ -794,7 +794,7 @@ def test_data_in_root(self): assert dt.children == {} assert_identical(dt.to_dataset(), dat) - def test_one_layer(self): + def test_one_layer(self) -> None: dat1, dat2 = xr.Dataset({"a": 1}), xr.Dataset({"b": 2}) dt = DataTree.from_dict({"run1": dat1, "run2": dat2}) assert_identical(dt.to_dataset(), xr.Dataset()) @@ -804,7 +804,7 @@ def test_one_layer(self): assert_identical(dt["run2"].to_dataset(), dat2) assert dt["run2"].children == {} - def test_two_layers(self): + def test_two_layers(self) -> None: dat1, dat2 = xr.Dataset({"a": 1}), xr.Dataset({"a": [1, 2]}) dt = DataTree.from_dict({"highres/run": dat1, "lowres/run": dat2}) assert "highres" in dt.children @@ -812,13 +812,13 @@ def test_two_layers(self): highres_run = dt["highres/run"] assert_identical(highres_run.to_dataset(), dat1) - def test_nones(self): + def test_nones(self) -> None: dt = DataTree.from_dict({"d": None, "d/e": None}) assert [node.name for node in dt.subtree] == [None, "d", "e"] assert [node.path for node in dt.subtree] == ["/", "/d", "/d/e"] assert_identical(dt["d/e"].to_dataset(), xr.Dataset()) - def test_full(self, simple_datatree): + def test_full(self, simple_datatree) -> None: dt = simple_datatree paths = list(node.path for node in dt.subtree) assert paths == [ @@ -831,7 +831,7 @@ def test_full(self, simple_datatree): "/set2/set1", ] - def test_datatree_values(self): + def test_datatree_values(self) -> None: dat1 = DataTree(dataset=xr.Dataset({"a": 1})) expected = DataTree() expected["a"] = dat1 @@ -840,13 +840,13 @@ def test_datatree_values(self): assert_identical(actual, expected) - def test_roundtrip(self, simple_datatree): + def test_roundtrip(self, simple_datatree) -> None: dt = simple_datatree roundtrip = DataTree.from_dict(dt.to_dict()) assert roundtrip.equals(dt) @pytest.mark.xfail - def test_roundtrip_unnamed_root(self, simple_datatree): + def test_roundtrip_unnamed_root(self, simple_datatree) -> None: # See GH81 dt = simple_datatree @@ -854,7 +854,7 @@ def test_roundtrip_unnamed_root(self, simple_datatree): roundtrip = DataTree.from_dict(dt.to_dict()) assert roundtrip.equals(dt) - def test_insertion_order(self): + def test_insertion_order(self) -> None: # regression test for GH issue #9276 reversed = DataTree.from_dict( { @@ -878,14 +878,14 @@ def test_insertion_order(self): # despite 'Bart' coming before 'Lisa' when sorted alphabetically assert list(reversed["Homer"].children.keys()) == ["Lisa", "Bart"] - def test_array_values(self): + def test_array_values(self) -> None: data = {"foo": xr.DataArray(1, name="bar")} with pytest.raises(TypeError): DataTree.from_dict(data) # type: ignore[arg-type] class TestDatasetView: - def test_view_contents(self): + def test_view_contents(self) -> None: ds = create_test_data() dt = DataTree(dataset=ds) assert ds.identical( @@ -893,7 +893,7 @@ def test_view_contents(self): ) # this only works because Dataset.identical doesn't check types assert isinstance(dt.dataset, xr.Dataset) - def test_immutability(self): + def test_immutability(self) -> None: # See issue https://github.com/xarray-contrib/datatree/issues/38 dt = DataTree.from_dict( { @@ -916,13 +916,13 @@ def test_immutability(self): # TODO are there any other ways you can normally modify state (in-place)? # (not attribute-like assignment because that doesn't work on Dataset anyway) - def test_methods(self): + def test_methods(self) -> None: ds = create_test_data() dt = DataTree(dataset=ds) assert ds.mean().identical(dt.dataset.mean()) assert isinstance(dt.dataset.mean(), xr.Dataset) - def test_arithmetic(self, create_test_datatree): + def test_arithmetic(self, create_test_datatree) -> None: dt = create_test_datatree() expected = create_test_datatree(modify=lambda ds: 10.0 * ds)[ "set1" @@ -930,7 +930,7 @@ def test_arithmetic(self, create_test_datatree): result = 10.0 * dt["set1"].dataset assert result.identical(expected) - def test_init_via_type(self): + def test_init_via_type(self) -> None: # from datatree GH issue https://github.com/xarray-contrib/datatree/issues/188 # xarray's .weighted is unusual because it uses type() to create a Dataset/DataArray @@ -948,7 +948,7 @@ def weighted_mean(ds): class TestAccess: - def test_attribute_access(self, create_test_datatree): + def test_attribute_access(self, create_test_datatree) -> None: dt = create_test_datatree() # vars / coords @@ -970,7 +970,7 @@ def test_attribute_access(self, create_test_datatree): assert dt.attrs["meta"] == "NASA" assert "meta" in dir(dt) - def test_ipython_key_completions(self, create_test_datatree): + def test_ipython_key_completions(self, create_test_datatree) -> None: dt = create_test_datatree() key_completions = dt._ipython_key_completions_() @@ -980,7 +980,7 @@ def test_ipython_key_completions(self, create_test_datatree): var_keys = list(dt.variables.keys()) assert all(var_key in key_completions for var_key in var_keys) - def test_operation_with_attrs_but_no_data(self): + def test_operation_with_attrs_but_no_data(self) -> None: # tests bug from xarray-datatree GH262 xs = xr.Dataset({"testvar": xr.DataArray(np.ones((2, 3)))}) dt = DataTree.from_dict({"node1": xs, "node2": xs}) @@ -990,7 +990,7 @@ def test_operation_with_attrs_but_no_data(self): class TestRepr: - def test_repr_four_nodes(self): + def test_repr_four_nodes(self) -> None: dt = DataTree.from_dict( { "/": xr.Dataset( @@ -1062,7 +1062,7 @@ def test_repr_four_nodes(self): ).strip() assert result == expected - def test_repr_two_children(self): + def test_repr_two_children(self) -> None: tree = DataTree.from_dict( { "/": Dataset(coords={"x": [1.0]}), @@ -1118,7 +1118,7 @@ def test_repr_two_children(self): ).strip() assert result == expected - def test_repr_inherited_dims(self): + def test_repr_inherited_dims(self) -> None: tree = DataTree.from_dict( { "/": Dataset({"foo": ("x", [1.0])}), @@ -1160,7 +1160,7 @@ def test_repr_inherited_dims(self): @pytest.mark.skipif( ON_WINDOWS, reason="windows (pre NumPy2) uses int32 instead of int64" ) - def test_doc_example(self): + def test_doc_example(self) -> None: # regression test for https://github.com/pydata/xarray/issues/9499 time = xr.DataArray(data=["2022-01", "2023-01"], dims="time") stations = xr.DataArray(data=list("abcdef"), dims="station") @@ -1262,7 +1262,7 @@ def _exact_match(message: str) -> str: class TestInheritance: - def test_inherited_dims(self): + def test_inherited_dims(self) -> None: dt = DataTree.from_dict( { "/": xr.Dataset({"d": (("x",), [1, 2])}), @@ -1279,7 +1279,7 @@ def test_inherited_dims(self): assert dt.b.to_dataset(inherit=True).sizes == {"y": 1} assert dt.b.to_dataset(inherit=False).sizes == {"y": 1} - def test_inherited_coords_index(self): + def test_inherited_coords_index(self) -> None: dt = DataTree.from_dict( { "/": xr.Dataset({"d": (("x",), [1, 2])}, coords={"x": [2, 3]}), @@ -1290,7 +1290,7 @@ def test_inherited_coords_index(self): assert "x" in dt["/b"].coords xr.testing.assert_identical(dt["/x"], dt["/b/x"]) - def test_inherit_only_index_coords(self): + def test_inherit_only_index_coords(self) -> None: dt = DataTree.from_dict( { "/": xr.Dataset(coords={"x": [1], "y": 2}), @@ -1308,7 +1308,7 @@ def test_inherit_only_index_coords(self): ) xr.testing.assert_equal(dt["/b/z"], xr.DataArray(3, coords={"z": 3})) - def test_inherited_coords_with_index_are_deduplicated(self): + def test_inherited_coords_with_index_are_deduplicated(self) -> None: dt = DataTree.from_dict( { "/": xr.Dataset(coords={"x": [1, 2]}), @@ -1324,7 +1324,7 @@ def test_inherited_coords_with_index_are_deduplicated(self): expected = xr.Dataset({"foo": ("x", [4, 5])}) assert_identical(child_dataset, expected) - def test_deduplicated_after_setitem(self): + def test_deduplicated_after_setitem(self) -> None: # regression test for GH #9601 dt = DataTree.from_dict( { @@ -1337,7 +1337,7 @@ def test_deduplicated_after_setitem(self): expected = xr.Dataset() assert_identical(child_dataset, expected) - def test_inconsistent_dims(self): + def test_inconsistent_dims(self) -> None: expected_msg = _exact_match( """ group '/b' is not aligned with its parents: @@ -1372,7 +1372,7 @@ def test_inconsistent_dims(self): children={"b": b}, ) - def test_inconsistent_child_indexes(self): + def test_inconsistent_child_indexes(self) -> None: expected_msg = _exact_match( """ group '/b' is not aligned with its parents: @@ -1407,7 +1407,7 @@ def test_inconsistent_child_indexes(self): with pytest.raises(ValueError, match=expected_msg): DataTree(dataset=xr.Dataset(coords={"x": [1.0]}), children={"b": b}) - def test_inconsistent_grandchild_indexes(self): + def test_inconsistent_grandchild_indexes(self) -> None: expected_msg = _exact_match( """ group '/b/c' is not aligned with its parents: @@ -1443,7 +1443,7 @@ def test_inconsistent_grandchild_indexes(self): with pytest.raises(ValueError, match=expected_msg): DataTree(dataset=xr.Dataset(coords={"x": [1.0]}), children={"b": b}) - def test_inconsistent_grandchild_dims(self): + def test_inconsistent_grandchild_dims(self) -> None: expected_msg = _exact_match( """ group '/b/c' is not aligned with its parents: @@ -1473,7 +1473,7 @@ def test_inconsistent_grandchild_dims(self): class TestRestructuring: - def test_drop_nodes(self): + def test_drop_nodes(self) -> None: sue = DataTree.from_dict({"Mary": None, "Kate": None, "Ashley": None}) # test drop just one node @@ -1493,7 +1493,7 @@ def test_drop_nodes(self): childless = dropped.drop_nodes(names=["Mary", "Ashley"], errors="ignore") assert childless.children == {} - def test_assign(self): + def test_assign(self) -> None: dt = DataTree() expected = DataTree.from_dict({"/": xr.Dataset({"foo": 0}), "/a": None}) @@ -1507,13 +1507,13 @@ def test_assign(self): class TestPipe: - def test_noop(self, create_test_datatree): + def test_noop(self, create_test_datatree) -> None: dt = create_test_datatree() actual = dt.pipe(lambda tree: tree) assert actual.identical(dt) - def test_params(self, create_test_datatree): + def test_params(self, create_test_datatree) -> None: dt = create_test_datatree() def f(tree, **attrs): @@ -1524,7 +1524,7 @@ def f(tree, **attrs): actual = dt.pipe(f, **attrs) assert actual["arr_with_attrs"].attrs == attrs - def test_named_self(self, create_test_datatree): + def test_named_self(self, create_test_datatree) -> None: dt = create_test_datatree() def f(x, tree, y): @@ -1539,7 +1539,7 @@ def f(x, tree, y): class TestSubset: - def test_match(self): + def test_match(self) -> None: # TODO is this example going to cause problems with case sensitivity? dt = DataTree.from_dict( { @@ -1558,7 +1558,7 @@ def test_match(self): ) assert_identical(result, expected) - def test_filter(self): + def test_filter(self) -> None: simpsons = DataTree.from_dict( { "/": xr.Dataset({"age": 83}), From 8d3137afd9bd7a392e16c423653f3b21561aaa8d Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 15 Oct 2024 21:04:10 -0400 Subject: [PATCH 5/5] more type hints --- xarray/tests/test_datatree.py | 78 +++++++++++++++++------------------ 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 692f57cc252..9c11cde3bbb 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -21,14 +21,14 @@ class TestTreeCreation: - def test_empty(self): + def test_empty(self) -> None: dt = DataTree(name="root") assert dt.name == "root" assert dt.parent is None assert dt.children == {} assert_identical(dt.to_dataset(), xr.Dataset()) - def test_name(self): + def test_name(self) -> None: dt = DataTree() assert dt.name is None @@ -50,14 +50,14 @@ def test_name(self): detached.name = "bar" assert detached.name == "bar" - def test_bad_names(self): + def test_bad_names(self) -> None: with pytest.raises(TypeError): DataTree(name=5) # type: ignore[arg-type] with pytest.raises(ValueError): DataTree(name="folder/data") - def test_data_arg(self): + def test_data_arg(self) -> None: ds = xr.Dataset({"foo": 42}) tree: DataTree = DataTree(dataset=ds) assert_identical(tree.to_dataset(), ds) @@ -67,13 +67,13 @@ def test_data_arg(self): class TestFamilyTree: - def test_dont_modify_children_inplace(self): + def test_dont_modify_children_inplace(self) -> None: # GH issue 9196 child = DataTree() DataTree(children={"child": child}) assert child.parent is None - def test_create_two_children(self): + def test_create_two_children(self) -> None: root_data = xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])}) set1_data = xr.Dataset({"a": 0, "b": 1}) root = DataTree.from_dict( @@ -82,7 +82,7 @@ def test_create_two_children(self): assert root["/set1"].name == "set1" assert root["/set1/set2"].name == "set2" - def test_create_full_tree(self, simple_datatree): + def test_create_full_tree(self, simple_datatree) -> None: d = simple_datatree.to_dict() d_keys = list(d.keys()) @@ -100,12 +100,12 @@ def test_create_full_tree(self, simple_datatree): class TestNames: - def test_child_gets_named_on_attach(self): + def test_child_gets_named_on_attach(self) -> None: sue = DataTree() mary = DataTree(children={"Sue": sue}) # noqa assert mary.children["Sue"].name == "Sue" - def test_dataset_containing_slashes(self): + def test_dataset_containing_slashes(self) -> None: xda: xr.DataArray = xr.DataArray( [[1, 2]], coords={"label": ["a"], "R30m/y": [30, 60]}, @@ -124,7 +124,7 @@ def test_dataset_containing_slashes(self): class TestPaths: - def test_path_property(self): + def test_path_property(self) -> None: john = DataTree.from_dict( { "/Mary/Sue": DataTree(), @@ -133,7 +133,7 @@ def test_path_property(self): assert john["/Mary/Sue"].path == "/Mary/Sue" assert john.path == "/" - def test_path_roundtrip(self): + def test_path_roundtrip(self) -> None: john = DataTree.from_dict( { "/Mary/Sue": DataTree(), @@ -141,7 +141,7 @@ def test_path_roundtrip(self): ) assert john["/Mary/Sue"].name == "Sue" - def test_same_tree(self): + def test_same_tree(self) -> None: john = DataTree.from_dict( { "/Mary": DataTree(), @@ -150,7 +150,7 @@ def test_same_tree(self): ) assert john["/Mary"].same_tree(john["/Kate"]) - def test_relative_paths(self): + def test_relative_paths(self) -> None: john = DataTree.from_dict( { "/Mary/Sue": DataTree(), @@ -179,7 +179,7 @@ def test_relative_paths(self): class TestStoreDatasets: - def test_create_with_data(self): + def test_create_with_data(self) -> None: dat = xr.Dataset({"a": 0}) john = DataTree(name="john", dataset=dat) @@ -188,7 +188,7 @@ def test_create_with_data(self): with pytest.raises(TypeError): DataTree(name="mary", dataset="junk") # type: ignore[arg-type] - def test_set_data(self): + def test_set_data(self) -> None: john = DataTree(name="john") dat = xr.Dataset({"a": 0}) john.dataset = dat # type: ignore[assignment] @@ -198,14 +198,14 @@ def test_set_data(self): with pytest.raises(TypeError): john.dataset = "junk" # type: ignore[assignment] - def test_has_data(self): + def test_has_data(self) -> None: john = DataTree(name="john", dataset=xr.Dataset({"a": 0})) assert john.has_data john_no_data = DataTree(name="john", dataset=None) assert not john_no_data.has_data - def test_is_hollow(self): + def test_is_hollow(self) -> None: john = DataTree(dataset=xr.Dataset({"a": 0})) assert john.is_hollow @@ -217,7 +217,7 @@ def test_is_hollow(self): class TestToDataset: - def test_to_dataset_inherited(self): + def test_to_dataset_inherited(self) -> None: base = xr.Dataset(coords={"a": [1], "b": 2}) sub = xr.Dataset(coords={"c": [3]}) tree = DataTree.from_dict({"/": base, "/sub": sub}) @@ -232,16 +232,16 @@ def test_to_dataset_inherited(self): class TestVariablesChildrenNameCollisions: - def test_parent_already_has_variable_with_childs_name(self): + def test_parent_already_has_variable_with_childs_name(self) -> None: with pytest.raises(KeyError, match="already contains a variable named a"): DataTree.from_dict({"/": xr.Dataset({"a": [0], "b": 1}), "/a": None}) - def test_parent_already_has_variable_with_childs_name_update(self): + def test_parent_already_has_variable_with_childs_name_update(self) -> None: dt = DataTree(dataset=xr.Dataset({"a": [0], "b": 1})) with pytest.raises(ValueError, match="already contains a variable named a"): dt.update({"a": DataTree()}) - def test_assign_when_already_child_with_variables_name(self): + def test_assign_when_already_child_with_variables_name(self) -> None: dt = DataTree.from_dict( { "/a": DataTree(), @@ -262,7 +262,7 @@ class TestGet: ... class TestGetItem: - def test_getitem_node(self): + def test_getitem_node(self) -> None: folder1 = DataTree.from_dict( { "/results/highres": DataTree(), @@ -272,16 +272,16 @@ def test_getitem_node(self): assert folder1["results"].name == "results" assert folder1["results/highres"].name == "highres" - def test_getitem_self(self): + def test_getitem_self(self) -> None: dt = DataTree() assert dt["."] is dt - def test_getitem_single_data_variable(self): + def test_getitem_single_data_variable(self) -> None: data = xr.Dataset({"temp": [0, 50]}) results = DataTree(name="results", dataset=data) assert_identical(results["temp"], data["temp"]) - def test_getitem_single_data_variable_from_node(self): + def test_getitem_single_data_variable_from_node(self) -> None: data = xr.Dataset({"temp": [0, 50]}) folder1 = DataTree.from_dict( { @@ -290,19 +290,19 @@ def test_getitem_single_data_variable_from_node(self): ) assert_identical(folder1["results/highres/temp"], data["temp"]) - def test_getitem_nonexistent_node(self): + def test_getitem_nonexistent_node(self) -> None: folder1 = DataTree.from_dict({"/results": DataTree()}, name="folder1") with pytest.raises(KeyError): folder1["results/highres"] - def test_getitem_nonexistent_variable(self): + def test_getitem_nonexistent_variable(self) -> None: data = xr.Dataset({"temp": [0, 50]}) results = DataTree(name="results", dataset=data) with pytest.raises(KeyError): results["pressure"] @pytest.mark.xfail(reason="Should be deprecated in favour of .subset") - def test_getitem_multiple_data_variables(self): + def test_getitem_multiple_data_variables(self) -> None: data = xr.Dataset({"temp": [0, 50], "p": [5, 8, 7]}) results = DataTree(name="results", dataset=data) assert_identical(results[["temp", "p"]], data[["temp", "p"]]) # type: ignore[index] @@ -310,47 +310,47 @@ def test_getitem_multiple_data_variables(self): @pytest.mark.xfail( reason="Indexing needs to return whole tree (GH https://github.com/xarray-contrib/datatree/issues/77)" ) - def test_getitem_dict_like_selection_access_to_dataset(self): + def test_getitem_dict_like_selection_access_to_dataset(self) -> None: data = xr.Dataset({"temp": [0, 50]}) results = DataTree(name="results", dataset=data) assert_identical(results[{"temp": 1}], data[{"temp": 1}]) # type: ignore[index] class TestUpdate: - def test_update(self): + def test_update(self) -> None: dt = DataTree() dt.update({"foo": xr.DataArray(0), "a": DataTree()}) expected = DataTree.from_dict({"/": xr.Dataset({"foo": 0}), "a": None}) assert_equal(dt, expected) assert dt.groups == ("/", "/a") - def test_update_new_named_dataarray(self): + def test_update_new_named_dataarray(self) -> None: da = xr.DataArray(name="temp", data=[0, 50]) folder1 = DataTree(name="folder1") folder1.update({"results": da}) expected = da.rename("results") assert_equal(folder1["results"], expected) - def test_update_doesnt_alter_child_name(self): + def test_update_doesnt_alter_child_name(self) -> None: dt = DataTree() dt.update({"foo": xr.DataArray(0), "a": DataTree(name="b")}) assert "a" in dt.children child = dt["a"] assert child.name == "a" - def test_update_overwrite(self): + def test_update_overwrite(self) -> None: actual = DataTree.from_dict({"a": DataTree(xr.Dataset({"x": 1}))}) actual.update({"a": DataTree(xr.Dataset({"x": 2}))}) expected = DataTree.from_dict({"a": DataTree(xr.Dataset({"x": 2}))}) assert_equal(actual, expected) - def test_update_coordinates(self): + def test_update_coordinates(self) -> None: expected = DataTree.from_dict({"/": xr.Dataset(coords={"a": 1})}) actual = DataTree.from_dict({"/": xr.Dataset()}) actual.update(xr.Dataset(coords={"a": 1})) assert_equal(actual, expected) - def test_update_inherited_coords(self): + def test_update_inherited_coords(self) -> None: expected = DataTree.from_dict( { "/": xr.Dataset(coords={"a": 1}), @@ -375,7 +375,7 @@ def test_update_inherited_coords(self): class TestCopy: - def test_copy(self, create_test_datatree): + def test_copy(self, create_test_datatree) -> None: dt = create_test_datatree() for node in dt.root.subtree: @@ -402,7 +402,7 @@ def test_copy(self, create_test_datatree): assert "foo" not in node.attrs assert node.attrs["Test"] is copied_node.attrs["Test"] - def test_copy_subtree(self): + def test_copy_subtree(self) -> None: dt = DataTree.from_dict({"/level1/level2/level3": xr.Dataset()}) actual = dt["/level1/level2"].copy() @@ -426,7 +426,7 @@ def test_copy_coord_inheritance(self) -> None: expected = DataTree(name="c") assert_identical(expected, actual) - def test_deepcopy(self, create_test_datatree): + def test_deepcopy(self, create_test_datatree) -> None: dt = create_test_datatree() for node in dt.root.subtree: @@ -454,7 +454,7 @@ def test_deepcopy(self, create_test_datatree): assert node.attrs["Test"] is not copied_node.attrs["Test"] @pytest.mark.xfail(reason="data argument not yet implemented") - def test_copy_with_data(self, create_test_datatree): + def test_copy_with_data(self, create_test_datatree) -> None: orig = create_test_datatree() # TODO use .data_vars once that property is available data_vars = {