From 065db0ee9166ae34a38971f91d988445115b8396 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Tue, 13 Feb 2024 16:02:46 -0700 Subject: [PATCH 01/19] Update the formating tests PR (#8702) added nbytes representation in DataArrays and Dataset repr, this adds it to the datatree tests. --- xarray/datatree_/datatree/tests/test_formatting.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/datatree_/datatree/tests/test_formatting.py b/xarray/datatree_/datatree/tests/test_formatting.py index 8726c95fe62..b58c02282e7 100644 --- a/xarray/datatree_/datatree/tests/test_formatting.py +++ b/xarray/datatree_/datatree/tests/test_formatting.py @@ -108,13 +108,13 @@ def test_diff_node_data(self): Data in nodes at position '/a' do not match: Data variables only on the left object: - v int64 1 + v int64 8B 1 Data in nodes at position '/a/b' do not match: Differing data variables: - L w int64 5 - R w int64 6""" + L w int64 8B 5 + R w int64 8B 6""" ) actual = diff_tree_repr(dt_1, dt_2, "equals") assert actual == expected From f167b99a29ebd4d397691a5695c8a7fd01cc93a8 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Tue, 13 Feb 2024 17:06:19 -0700 Subject: [PATCH 02/19] Migrate treenode module Moves treenode.py and test_treenode.py. Updates some typing. Updates imports from treenode. --- xarray/backends/common.py | 4 +- xarray/backends/zarr.py | 4 +- .../{datatree_/datatree => core}/treenode.py | 40 +++++++++---------- xarray/datatree_/datatree/__init__.py | 2 +- xarray/datatree_/datatree/datatree.py | 2 +- xarray/datatree_/datatree/iterators.py | 2 +- xarray/datatree_/datatree/mapping.py | 4 +- .../tests => tests/datatree}/test_treenode.py | 4 +- 8 files changed, 31 insertions(+), 31 deletions(-) rename xarray/{datatree_/datatree => core}/treenode.py (95%) rename xarray/{datatree_/datatree/tests => tests/datatree}/test_treenode.py (98%) diff --git a/xarray/backends/common.py b/xarray/backends/common.py index 6245b3442a3..7d3cc00a52d 100644 --- a/xarray/backends/common.py +++ b/xarray/backends/common.py @@ -137,8 +137,8 @@ def _open_datatree_netcdf( **kwargs, ) -> DataTree: from xarray.backends.api import open_dataset + from xarray.core.treenode import NodePath from xarray.datatree_.datatree import DataTree - from xarray.datatree_.datatree.treenode import NodePath ds = open_dataset(filename_or_obj, **kwargs) tree_root = DataTree.from_dict({"/": ds}) @@ -159,7 +159,7 @@ def _open_datatree_netcdf( def _iter_nc_groups(root, parent="/"): - from xarray.datatree_.datatree.treenode import NodePath + from xarray.core.treenode import NodePath parent = NodePath(parent) for path, group in root.groups.items(): diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index ac208da097a..e9465dc0ba0 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -1048,8 +1048,8 @@ def open_datatree( import zarr from xarray.backends.api import open_dataset + from xarray.core.treenode import NodePath from xarray.datatree_.datatree import DataTree - from xarray.datatree_.datatree.treenode import NodePath zds = zarr.open_group(filename_or_obj, mode="r") ds = open_dataset(filename_or_obj, engine="zarr", **kwargs) @@ -1075,7 +1075,7 @@ def open_datatree( def _iter_zarr_groups(root, parent="/"): - from xarray.datatree_.datatree.treenode import NodePath + from xarray.core.treenode import NodePath parent = NodePath(parent) for path, group in root.groups(): diff --git a/xarray/datatree_/datatree/treenode.py b/xarray/core/treenode.py similarity index 95% rename from xarray/datatree_/datatree/treenode.py rename to xarray/core/treenode.py index 1689d261c34..c9e8d4335e4 100644 --- a/xarray/datatree_/datatree/treenode.py +++ b/xarray/core/treenode.py @@ -2,16 +2,12 @@ import sys from collections import OrderedDict +from collections.abc import Iterator, Mapping from pathlib import PurePosixPath from typing import ( TYPE_CHECKING, Generic, - Iterator, - Mapping, - Optional, - Tuple, TypeVar, - Union, ) from xarray.core.utils import Frozen, is_dict_like @@ -24,6 +20,8 @@ class InvalidTreeError(Exception): """Raised when user attempts to create an invalid tree in some way.""" +# TODO [MHS, 02/13/2024] I don't like the description here. It doesn't make +# sense on first glance. class NotFoundInTreeError(ValueError): """Raised when operation can't be completed because one node is part of the expected tree.""" @@ -75,10 +73,10 @@ class TreeNode(Generic[Tree]): (This class is heavily inspired by the anytree library's NodeMixin class.) """ - _parent: Optional[Tree] + _parent: Tree | None _children: OrderedDict[str, Tree] - def __init__(self, children: Optional[Mapping[str, Tree]] = None): + def __init__(self, children: Mapping[str, Tree] | None = None): """Create a parentless node.""" self._parent = None self._children = OrderedDict() @@ -91,7 +89,7 @@ def parent(self) -> Tree | None: return self._parent def _set_parent( - self, new_parent: Tree | None, child_name: Optional[str] = None + self, new_parent: Tree | None, child_name: str | None = None ) -> None: # TODO is it possible to refactor in a way that removes this private method? @@ -137,7 +135,7 @@ def _detach(self, parent: Tree | None) -> None: self._parent = None self._post_detach(parent) - def _attach(self, parent: Tree | None, child_name: Optional[str] = None) -> None: + def _attach(self, parent: Tree | None, child_name: str | None = None) -> None: if parent is not None: if child_name is None: raise ValueError( @@ -242,7 +240,7 @@ def _iter_parents(self: Tree) -> Iterator[Tree]: yield node node = node.parent - def iter_lineage(self: Tree) -> Tuple[Tree, ...]: + def iter_lineage(self: Tree) -> tuple[Tree, ...]: """Iterate up the tree, starting from the current node.""" from warnings import warn @@ -254,7 +252,7 @@ def iter_lineage(self: Tree) -> Tuple[Tree, ...]: return tuple((self, *self.parents)) @property - def lineage(self: Tree) -> Tuple[Tree, ...]: + def lineage(self: Tree) -> tuple[Tree, ...]: """All parent nodes and their parent nodes, starting with the closest.""" from warnings import warn @@ -266,12 +264,12 @@ def lineage(self: Tree) -> Tuple[Tree, ...]: return self.iter_lineage() @property - def parents(self: Tree) -> Tuple[Tree, ...]: + def parents(self: Tree) -> tuple[Tree, ...]: """All parent nodes and their parent nodes, starting with the closest.""" return tuple(self._iter_parents()) @property - def ancestors(self: Tree) -> Tuple[Tree, ...]: + def ancestors(self: Tree) -> tuple[Tree, ...]: """All parent nodes and their parent nodes, starting with the most distant.""" from warnings import warn @@ -306,7 +304,7 @@ def is_leaf(self) -> bool: return self.children == {} @property - def leaves(self: Tree) -> Tuple[Tree, ...]: + def leaves(self: Tree) -> tuple[Tree, ...]: """ All leaf nodes. @@ -341,12 +339,12 @@ def subtree(self: Tree) -> Iterator[Tree]: -------- DataTree.descendants """ - from . import iterators + from xarray.datatree_.datatree import iterators return iterators.PreOrderIter(self) @property - def descendants(self: Tree) -> Tuple[Tree, ...]: + def descendants(self: Tree) -> tuple[Tree, ...]: """ Child nodes and all their child nodes. @@ -431,7 +429,7 @@ def _post_attach(self: Tree, parent: Tree) -> None: """Method call after attaching to `parent`.""" pass - def get(self: Tree, key: str, default: Optional[Tree] = None) -> Optional[Tree]: + def get(self: Tree, key: str, default: Tree | None = None) -> Tree | None: """ Return the child node with the specified key. @@ -445,7 +443,7 @@ def get(self: Tree, key: str, default: Optional[Tree] = None) -> Optional[Tree]: # TODO `._walk` method to be called by both `_get_item` and `_set_item` - def _get_item(self: Tree, path: str | NodePath) -> Union[Tree, T_DataArray]: + def _get_item(self: Tree, path: str | NodePath) -> Tree | T_DataArray: """ Returns the object lying at the given path. @@ -488,7 +486,7 @@ def _set(self: Tree, key: str, val: Tree) -> None: def _set_item( self: Tree, path: str | NodePath, - item: Union[Tree, T_DataArray], + item: Tree | T_DataArray, new_nodes_along_path: bool = False, allow_overwrite: bool = True, ) -> None: @@ -580,8 +578,8 @@ class NamedNode(TreeNode, Generic[Tree]): Implements path-like relationships to other nodes in its tree. """ - _name: Optional[str] - _parent: Optional[Tree] + _name: str | None + _parent: Tree | None _children: OrderedDict[str, Tree] def __init__(self, name=None, children=None): diff --git a/xarray/datatree_/datatree/__init__.py b/xarray/datatree_/datatree/__init__.py index f9fd419bddc..071dcbecf8c 100644 --- a/xarray/datatree_/datatree/__init__.py +++ b/xarray/datatree_/datatree/__init__.py @@ -2,7 +2,7 @@ from .datatree import DataTree from .extensions import register_datatree_accessor from .mapping import TreeIsomorphismError, map_over_subtree -from .treenode import InvalidTreeError, NotFoundInTreeError +from xarray.core.treenode import InvalidTreeError, NotFoundInTreeError __all__ = ( diff --git a/xarray/datatree_/datatree/datatree.py b/xarray/datatree_/datatree/datatree.py index 13cca7de80d..10133052185 100644 --- a/xarray/datatree_/datatree/datatree.py +++ b/xarray/datatree_/datatree/datatree.py @@ -50,7 +50,7 @@ MappedDataWithCoords, ) from .render import RenderTree -from .treenode import NamedNode, NodePath, Tree +from xarray.core.treenode import NamedNode, NodePath, Tree try: from xarray.core.variable import calculate_dimensions diff --git a/xarray/datatree_/datatree/iterators.py b/xarray/datatree_/datatree/iterators.py index 52ed8d22422..68e75c4f612 100644 --- a/xarray/datatree_/datatree/iterators.py +++ b/xarray/datatree_/datatree/iterators.py @@ -2,7 +2,7 @@ from collections import abc from typing import Callable, Iterator, List, Optional -from .treenode import Tree +from xarray.core.treenode import Tree """These iterators are copied from anytree.iterators, with minor modifications.""" diff --git a/xarray/datatree_/datatree/mapping.py b/xarray/datatree_/datatree/mapping.py index 34e227d349d..355149060a9 100644 --- a/xarray/datatree_/datatree/mapping.py +++ b/xarray/datatree_/datatree/mapping.py @@ -9,10 +9,10 @@ from xarray import DataArray, Dataset from .iterators import LevelOrderIter -from .treenode import NodePath, TreeNode +from xarray.core.treenode import NodePath, TreeNode if TYPE_CHECKING: - from .datatree import DataTree + from xarray.core.datatree import DataTree class TreeIsomorphismError(ValueError): diff --git a/xarray/datatree_/datatree/tests/test_treenode.py b/xarray/tests/datatree/test_treenode.py similarity index 98% rename from xarray/datatree_/datatree/tests/test_treenode.py rename to xarray/tests/datatree/test_treenode.py index 3c75f3ac8a4..245e8fab4fc 100644 --- a/xarray/datatree_/datatree/tests/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -1,7 +1,9 @@ +from __future__ import annotations + import pytest +from xarray.core.treenode import InvalidTreeError, NamedNode, NodePath, TreeNode from xarray.datatree_.datatree.iterators import LevelOrderIter, PreOrderIter -from xarray.datatree_.datatree.treenode import InvalidTreeError, NamedNode, NodePath, TreeNode class TestFamilyTree: From 6bd492f6d3ac0a8b6082316af0ec0b10c90027a9 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Thu, 15 Feb 2024 09:04:00 -0700 Subject: [PATCH 03/19] Update NotFoundInTreeError description. --- xarray/core/treenode.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index c9e8d4335e4..019726c8520 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -20,10 +20,8 @@ class InvalidTreeError(Exception): """Raised when user attempts to create an invalid tree in some way.""" -# TODO [MHS, 02/13/2024] I don't like the description here. It doesn't make -# sense on first glance. class NotFoundInTreeError(ValueError): - """Raised when operation can't be completed because one node is part of the expected tree.""" + """Raised when operation can't be completed because one node is not part of the expected tree.""" class NodePath(PurePosixPath): From b62a21a67fd0906fb4765aee1041c89721ef89c2 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Thu, 15 Feb 2024 13:58:32 -0700 Subject: [PATCH 04/19] Reformat some comments Add test tree structure for easier understanding. --- xarray/core/treenode.py | 10 ++++++---- xarray/tests/datatree/test_treenode.py | 26 ++++++++++++++++++++------ 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 019726c8520..273c2c35a2d 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -491,17 +491,19 @@ def _set_item( """ Set a new item in the tree, overwriting anything already present at that path. - The given value either forms a new node of the tree or overwrites an existing item at that location. + The given value either forms a new node of the tree or overwrites an + existing item at that location. Parameters ---------- path item new_nodes_along_path : bool - If true, then if necessary new nodes will be created along the given path, until the tree can reach the - specified location. + If true, then if necessary new nodes will be created along the + given path, until the tree can reach the specified location. allow_overwrite : bool - Whether or not to overwrite any existing node at the location given by path. + Whether or not to overwrite any existing node at the location given + by path. Raises ------ diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/datatree/test_treenode.py index 245e8fab4fc..a7d3a17ec1f 100644 --- a/xarray/tests/datatree/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -227,6 +227,16 @@ def test_del_child(self): def create_test_tree(): + # a + # ├── b + # │ ├── d + # │ └── e + # │ ├── f + # │ └── g + # └── c + # └── h + # └── i + a = NamedNode(name="a") b = NamedNode() c = NamedNode() @@ -281,19 +291,20 @@ def test_levelorderiter(self): class TestAncestry: + def test_parents(self): - _, leaf = create_test_tree() + _, leaf_f = create_test_tree() expected = ["e", "b", "a"] - assert [node.name for node in leaf.parents] == expected + assert [node.name for node in leaf_f.parents] == expected def test_lineage(self): - _, leaf = create_test_tree() + _, leaf_f = create_test_tree() expected = ["f", "e", "b", "a"] - assert [node.name for node in leaf.lineage] == expected + assert [node.name for node in leaf_f.lineage] == expected def test_ancestors(self): - _, leaf = create_test_tree() - ancestors = leaf.ancestors + _, leaf_f = create_test_tree() + ancestors = leaf_f.ancestors expected = ["a", "b", "e", "f"] for node, expected_name in zip(ancestors, expected): assert node.name == expected_name @@ -372,6 +383,9 @@ def test_render_nodetree(self): "NamedNode('Ben')", "NamedNode('Kate')", ] + + john_str = printout.splitlines() + assert len(john_str) == len(expected_nodes) for expected_node, printed_node in zip(expected_nodes, printout.splitlines()): assert expected_node in printed_node From 32053b6796df5c895714591b23d438c53a2ae24e Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Thu, 15 Feb 2024 14:04:29 -0700 Subject: [PATCH 05/19] Updates whats-new.rst --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 1eca1d30e72..a32957461aa 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -102,6 +102,8 @@ Internal Changes - Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By `Matt Savoie `_. +- Adds ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By `Matt Savoie `_. + .. _whats-new.2024.01.1: v2024.01.1 (23 Jan, 2024) From 32e7453fdf65671845593f963058601eb30a2e4a Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Thu, 15 Feb 2024 14:16:07 -0700 Subject: [PATCH 06/19] mypy typing. (terrible?) There must be a better way, but I don't know it. particularly the list comprehension casts. --- xarray/tests/datatree/test_treenode.py | 124 +++++++++++++------------ 1 file changed, 65 insertions(+), 59 deletions(-) diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/datatree/test_treenode.py index a7d3a17ec1f..42a55fb4b68 100644 --- a/xarray/tests/datatree/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -1,5 +1,8 @@ from __future__ import annotations +from collections.abc import Iterator +from typing import cast + import pytest from xarray.core.treenode import InvalidTreeError, NamedNode, NodePath, TreeNode @@ -8,20 +11,20 @@ class TestFamilyTree: def test_lonely(self): - root = TreeNode() + root: TreeNode = TreeNode() assert root.parent is None assert root.children == {} def test_parenting(self): - john = TreeNode() - mary = TreeNode() + john: TreeNode = TreeNode() + mary: TreeNode = TreeNode() mary._set_parent(john, "Mary") assert mary.parent == john assert john.children["Mary"] is mary def test_no_time_traveller_loops(self): - john = TreeNode() + john: TreeNode = TreeNode() with pytest.raises(InvalidTreeError, match="cannot be a parent of itself"): john._set_parent(john, "John") @@ -29,8 +32,8 @@ def test_no_time_traveller_loops(self): with pytest.raises(InvalidTreeError, match="cannot be a parent of itself"): john.children = {"John": john} - mary = TreeNode() - rose = TreeNode() + mary: TreeNode = TreeNode() + rose: TreeNode = TreeNode() mary._set_parent(john, "Mary") rose._set_parent(mary, "Rose") @@ -41,11 +44,11 @@ def test_no_time_traveller_loops(self): rose.children = {"John": john} def test_parent_swap(self): - john = TreeNode() - mary = TreeNode() + john: TreeNode = TreeNode() + mary: TreeNode = TreeNode() mary._set_parent(john, "Mary") - steve = TreeNode() + steve: TreeNode = TreeNode() mary._set_parent(steve, "Mary") assert mary.parent == steve @@ -53,24 +56,24 @@ def test_parent_swap(self): assert "Mary" not in john.children def test_multi_child_family(self): - mary = TreeNode() - kate = TreeNode() - john = TreeNode(children={"Mary": mary, "Kate": kate}) + mary: TreeNode = TreeNode() + kate: TreeNode = TreeNode() + john: TreeNode = TreeNode(children={"Mary": mary, "Kate": kate}) assert john.children["Mary"] is mary assert john.children["Kate"] is kate assert mary.parent is john assert kate.parent is john def test_disown_child(self): - mary = TreeNode() - john = TreeNode(children={"Mary": mary}) + mary: TreeNode = TreeNode() + john: TreeNode = TreeNode(children={"Mary": mary}) mary.orphan() assert mary.parent is None assert "Mary" not in john.children def test_doppelganger_child(self): - kate = TreeNode() - john = TreeNode() + kate: TreeNode = TreeNode() + john: TreeNode = TreeNode() with pytest.raises(TypeError): john.children = {"Kate": 666} @@ -79,22 +82,22 @@ def test_doppelganger_child(self): john.children = {"Kate": kate, "Evil_Kate": kate} john = TreeNode(children={"Kate": kate}) - evil_kate = TreeNode() + evil_kate: TreeNode = TreeNode() evil_kate._set_parent(john, "Kate") assert john.children["Kate"] is evil_kate def test_sibling_relationships(self): - mary = TreeNode() - kate = TreeNode() - ashley = TreeNode() + mary: TreeNode = TreeNode() + kate: TreeNode = TreeNode() + ashley: TreeNode = TreeNode() TreeNode(children={"Mary": mary, "Kate": kate, "Ashley": ashley}) assert kate.siblings["Mary"] is mary assert kate.siblings["Ashley"] is ashley assert "Kate" not in kate.siblings def test_ancestors(self): - tony = TreeNode() - michael = TreeNode(children={"Tony": tony}) + tony: TreeNode = TreeNode() + michael: TreeNode = TreeNode(children={"Tony": tony}) vito = TreeNode(children={"Michael": michael}) assert tony.root is vito assert tony.parents == (michael, vito) @@ -103,7 +106,7 @@ def test_ancestors(self): class TestGetNodes: def test_get_child(self): - steven = TreeNode() + steven: TreeNode = TreeNode() sue = TreeNode(children={"Steven": steven}) mary = TreeNode(children={"Sue": sue}) john = TreeNode(children={"Mary": mary}) @@ -126,8 +129,8 @@ def test_get_child(self): assert mary._get_item("Sue/Steven") is steven def test_get_upwards(self): - sue = TreeNode() - kate = TreeNode() + sue: TreeNode = TreeNode() + kate: TreeNode = TreeNode() mary = TreeNode(children={"Sue": sue, "Kate": kate}) john = TreeNode(children={"Mary": mary}) @@ -138,7 +141,7 @@ def test_get_upwards(self): assert sue._get_item("../Kate") is kate def test_get_from_root(self): - sue = TreeNode() + sue: TreeNode = TreeNode() mary = TreeNode(children={"Sue": sue}) john = TreeNode(children={"Mary": mary}) # noqa @@ -147,8 +150,8 @@ def test_get_from_root(self): class TestSetNodes: def test_set_child_node(self): - john = TreeNode() - mary = TreeNode() + john: TreeNode = TreeNode() + mary: TreeNode = TreeNode() john._set_item("Mary", mary) assert john.children["Mary"] is mary @@ -157,16 +160,16 @@ def test_set_child_node(self): assert mary.parent is john def test_child_already_exists(self): - mary = TreeNode() - john = TreeNode(children={"Mary": mary}) - mary_2 = TreeNode() + mary: TreeNode = TreeNode() + john: TreeNode = TreeNode(children={"Mary": mary}) + mary_2: TreeNode = TreeNode() with pytest.raises(KeyError): john._set_item("Mary", mary_2, allow_overwrite=False) def test_set_grandchild(self): - rose = TreeNode() - mary = TreeNode() - john = TreeNode() + rose: TreeNode = TreeNode() + mary: TreeNode = TreeNode() + john: TreeNode = TreeNode() john._set_item("Mary", mary) john._set_item("Mary/Rose", rose) @@ -177,8 +180,8 @@ def test_set_grandchild(self): assert rose.parent is mary def test_create_intermediate_child(self): - john = TreeNode() - rose = TreeNode() + john: TreeNode = TreeNode() + rose: TreeNode = TreeNode() # test intermediate children not allowed with pytest.raises(KeyError, match="Could not reach"): @@ -194,12 +197,12 @@ def test_create_intermediate_child(self): assert rose.parent == mary def test_overwrite_child(self): - john = TreeNode() - mary = TreeNode() + john: TreeNode = TreeNode() + mary: TreeNode = TreeNode() john._set_item("Mary", mary) # test overwriting not allowed - marys_evil_twin = TreeNode() + marys_evil_twin: TreeNode = TreeNode() with pytest.raises(KeyError, match="Already a node object"): john._set_item("Mary", marys_evil_twin, allow_overwrite=False) assert john.children["Mary"] is mary @@ -214,8 +217,8 @@ def test_overwrite_child(self): class TestPruning: def test_del_child(self): - john = TreeNode() - mary = TreeNode() + john: TreeNode = TreeNode() + mary: TreeNode = TreeNode() john._set_item("Mary", mary) del john["Mary"] @@ -226,7 +229,7 @@ def test_del_child(self): del john["Mary"] -def create_test_tree(): +def create_test_tree() -> tuple[NamedNode, NamedNode]: # a # ├── b # │ ├── d @@ -236,16 +239,15 @@ def create_test_tree(): # └── c # └── h # └── i - - a = NamedNode(name="a") - b = NamedNode() - c = NamedNode() - d = NamedNode() - e = NamedNode() - f = NamedNode() - g = NamedNode() - h = NamedNode() - i = NamedNode() + a: NamedNode = NamedNode(name="a") + b: NamedNode = NamedNode() + c: NamedNode = NamedNode() + d: NamedNode = NamedNode() + e: NamedNode = NamedNode() + f: NamedNode = NamedNode() + g: NamedNode = NamedNode() + h: NamedNode = NamedNode() + i: NamedNode = NamedNode() a.children = {"b": b, "c": c} b.children = {"d": d, "e": e} @@ -259,7 +261,9 @@ def create_test_tree(): class TestIterators: def test_preorderiter(self): root, _ = create_test_tree() - result = [node.name for node in PreOrderIter(root)] + result: list[str | None] = [ + node.name for node in cast(Iterator[NamedNode], PreOrderIter(root)) + ] expected = [ "a", "b", @@ -275,7 +279,9 @@ def test_preorderiter(self): def test_levelorderiter(self): root, _ = create_test_tree() - result = [node.name for node in LevelOrderIter(root)] + result: list[str | None] = [ + node.name for node in cast(Iterator[NamedNode], LevelOrderIter(root)) + ] expected = [ "a", # root Node is unnamed "b", @@ -369,11 +375,11 @@ def test_levels(self): class TestRenderTree: def test_render_nodetree(self): - sam = NamedNode() - ben = NamedNode() - mary = NamedNode(children={"Sam": sam, "Ben": ben}) - kate = NamedNode() - john = NamedNode(children={"Mary": mary, "Kate": kate}) + sam: NamedNode = NamedNode() + ben: NamedNode = NamedNode() + mary: NamedNode = NamedNode(children={"Sam": sam, "Ben": ben}) + kate: NamedNode = NamedNode() + john: NamedNode = NamedNode(children={"Mary": mary, "Kate": kate}) printout = john.__str__() expected_nodes = [ From 7f2a178ece762ae2b67a9e3ffed1d2a2571ab15b Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Fri, 16 Feb 2024 10:57:40 -0700 Subject: [PATCH 07/19] Adds __repr__ to NamedNode and updates test This test was broken becuase only the root node was being tested and none of the previous nodes were represented in the __str__. --- xarray/core/treenode.py | 6 ++++++ xarray/tests/datatree/test_treenode.py | 24 ++++++++++++++---------- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 273c2c35a2d..341aada373d 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -601,6 +601,12 @@ def name(self, name: str | None) -> None: raise ValueError("node names cannot contain forward slashes") self._name = name + def __repr__(self, level=0): + repr_value = "\t" * level + self.__str__() + "\n" + for child in self.children: + repr_value += self.get(child).__repr__(level + 1) + return repr_value + def __str__(self) -> str: return f"NamedNode({self.name})" if self.name else "NamedNode()" diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/datatree/test_treenode.py index 42a55fb4b68..114454e0544 100644 --- a/xarray/tests/datatree/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -380,20 +380,24 @@ def test_render_nodetree(self): mary: NamedNode = NamedNode(children={"Sam": sam, "Ben": ben}) kate: NamedNode = NamedNode() john: NamedNode = NamedNode(children={"Mary": mary, "Kate": kate}) - - printout = john.__str__() expected_nodes = [ "NamedNode()", - "NamedNode('Mary')", - "NamedNode('Sam')", - "NamedNode('Ben')", - "NamedNode('Kate')", + "\tNamedNode(Mary)", + "\t\tNamedNode(Sam)", + "\t\tNamedNode(Ben)", + "\tNamedNode(Kate)", ] + expected_str = "NamedNode(Mary)" + + john_repr = john.__repr__() + mary_str = mary.__str__() + + assert mary_str == expected_str - john_str = printout.splitlines() - assert len(john_str) == len(expected_nodes) - for expected_node, printed_node in zip(expected_nodes, printout.splitlines()): - assert expected_node in printed_node + john_nodes = john_repr.splitlines() + assert len(john_nodes) == len(expected_nodes) + for expected_node, repr_node in zip(expected_nodes, john_nodes): + assert expected_node == repr_node def test_nodepath(): From 548536b0083c2bcfbd8d56bee9434a0113c0c229 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Fri, 16 Feb 2024 12:06:36 -0700 Subject: [PATCH 08/19] Adds quotes to NamedNode __str__ representation. --- xarray/core/treenode.py | 2 +- xarray/tests/datatree/test_treenode.py | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 341aada373d..ab4f0326cfc 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -608,7 +608,7 @@ def __repr__(self, level=0): return repr_value def __str__(self) -> str: - return f"NamedNode({self.name})" if self.name else "NamedNode()" + return f'NamedNode("{self.name}")' if self.name else "NamedNode()" def _post_attach(self: NamedNode, parent: NamedNode) -> None: """Ensures child has name attribute corresponding to key under which it has been stored.""" diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/datatree/test_treenode.py index 114454e0544..9435881bfef 100644 --- a/xarray/tests/datatree/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -382,13 +382,12 @@ def test_render_nodetree(self): john: NamedNode = NamedNode(children={"Mary": mary, "Kate": kate}) expected_nodes = [ "NamedNode()", - "\tNamedNode(Mary)", - "\t\tNamedNode(Sam)", - "\t\tNamedNode(Ben)", - "\tNamedNode(Kate)", + '\tNamedNode("Mary")', + '\t\tNamedNode("Sam")', + '\t\tNamedNode("Ben")', + '\tNamedNode("Kate")', ] - expected_str = "NamedNode(Mary)" - + expected_str = 'NamedNode("Mary")' john_repr = john.__repr__() mary_str = mary.__str__() From b9685d770ddee8466a150bf0eca3542e88f273b6 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Fri, 16 Feb 2024 12:13:29 -0700 Subject: [PATCH 09/19] swaps " for ' in NamedNode __str__ representation. --- xarray/core/treenode.py | 2 +- xarray/tests/datatree/test_treenode.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index ab4f0326cfc..c408c43d740 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -608,7 +608,7 @@ def __repr__(self, level=0): return repr_value def __str__(self) -> str: - return f'NamedNode("{self.name}")' if self.name else "NamedNode()" + return f"NamedNode('{self.name}')" if self.name else "NamedNode()" def _post_attach(self: NamedNode, parent: NamedNode) -> None: """Ensures child has name attribute corresponding to key under which it has been stored.""" diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/datatree/test_treenode.py index 9435881bfef..b0e737bd317 100644 --- a/xarray/tests/datatree/test_treenode.py +++ b/xarray/tests/datatree/test_treenode.py @@ -382,12 +382,12 @@ def test_render_nodetree(self): john: NamedNode = NamedNode(children={"Mary": mary, "Kate": kate}) expected_nodes = [ "NamedNode()", - '\tNamedNode("Mary")', - '\t\tNamedNode("Sam")', - '\t\tNamedNode("Ben")', - '\tNamedNode("Kate")', + "\tNamedNode('Mary')", + "\t\tNamedNode('Sam')", + "\t\tNamedNode('Ben')", + "\tNamedNode('Kate')", ] - expected_str = 'NamedNode("Mary")' + expected_str = "NamedNode('Mary')" john_repr = john.__repr__() mary_str = mary.__str__() From 59da654e153d1128a497ca9c5228d0f8d8f1d84f Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Fri, 16 Feb 2024 13:09:35 -0700 Subject: [PATCH 10/19] Adding Tom in so he gets blamed properly. --- doc/whats-new.rst | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index fd5515501c8..e1cf74ab113 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -96,13 +96,18 @@ Internal Changes By `Tom Nicholas `_ and `Anderson Banihirwe `_. - Imports ``datatree`` repository and history into internal - location. (:pull:`8688`) By `Matt Savoie `_ - and `Justus Magin `_. + location. (:pull:`8688`) By `Matt Savoie `_, + `Justus Magin `_ and `Tom Nicholas + `_. -- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By `Matt - Savoie `_. +- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By + `Matt Savoie `_ and `Tom Nicholas + `_. + +- Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By + `Matt Savoie `_ and `Tom Nicholas + `_. -- Adds ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By `Matt Savoie `_. - Refactor :py:meth:`xarray.core.indexing.DaskIndexingAdapter.__getitem__` to remove an unnecessary rewrite of the indexer key (:issue: `8377`, :pull:`8758`) By `Anderson Banihirwe ` From 4530973a58605b02ddcd6677ce540d6ed5ed6851 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Mon, 19 Feb 2024 12:39:04 -0700 Subject: [PATCH 11/19] resolve conflict whats-new.rst Question is I did update below the released line to give Tom some credit. I hope that's is allowable. --- doc/whats-new.rst | 44 +++++++++++++++++++++++++++++++------------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 7678a67437c..dd3e4d7d87f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -15,6 +15,37 @@ What's New np.random.seed(123456) +.. _whats-new.2024.03.0: + +v2024.03.0 (unreleased) +----------------------- + +New Features +~~~~~~~~~~~~ + + +Breaking changes +~~~~~~~~~~~~~~~~ + + +Deprecations +~~~~~~~~~~~~ + + +Bug fixes +~~~~~~~~~ + + +Documentation +~~~~~~~~~~~~~ + + +Internal Changes +~~~~~~~~~~~~~~~~ +- Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By + `Matt Savoie `_ and `Tom Nicholas + `_. + .. _whats-new.2024.02.0: v2024.02.0 (Feb 19, 2024) @@ -107,26 +138,13 @@ Internal Changes - Move ``parallelcompat`` and ``chunk managers`` modules from ``xarray/core`` to ``xarray/namedarray``. (:pull:`8319`) By `Tom Nicholas `_ and `Anderson Banihirwe `_. - - Imports ``datatree`` repository and history into internal location. (:pull:`8688`) By `Matt Savoie `_, `Justus Magin `_ and `Tom Nicholas `_. - - Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By `Matt Savoie `_ and `Tom Nicholas `_. - -- Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By - `Matt Savoie `_ and `Tom Nicholas - `_. - -- Refactor :py:meth:`xarray.core.indexing.DaskIndexingAdapter.__getitem__` to remove an unnecessary rewrite of the indexer key - (:issue: `8377`, :pull:`8758`) By `Anderson Banihirwe ` -- Imports ``datatree`` repository and history into internal location. (:pull:`8688`) - By `Matt Savoie `_ and `Justus Magin `_. -- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) - By `Matt Savoie `_. - Refactor :py:meth:`xarray.core.indexing.DaskIndexingAdapter.__getitem__` to remove an unnecessary rewrite of the indexer key (:issue: `8377`, :pull:`8758`) By `Anderson Banihirwe `_. From c03d373c915436e5e816a24975f9f686bae15653 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Mon, 19 Feb 2024 15:07:14 -0700 Subject: [PATCH 12/19] Moves test_treenode.py to xarray/tests. Integrated tests. --- doc/whats-new.rst | 15 +++++++-------- xarray/tests/{datatree => }/test_treenode.py | 0 2 files changed, 7 insertions(+), 8 deletions(-) rename xarray/tests/{datatree => }/test_treenode.py (100%) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index dd3e4d7d87f..f991f6bdfe6 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -42,8 +42,8 @@ Documentation Internal Changes ~~~~~~~~~~~~~~~~ -- Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By - `Matt Savoie `_ and `Tom Nicholas +- Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) + By `Matt Savoie `_ and `Tom Nicholas `_. .. _whats-new.2024.02.0: @@ -138,12 +138,11 @@ Internal Changes - Move ``parallelcompat`` and ``chunk managers`` modules from ``xarray/core`` to ``xarray/namedarray``. (:pull:`8319`) By `Tom Nicholas `_ and `Anderson Banihirwe `_. -- Imports ``datatree`` repository and history into internal - location. (:pull:`8688`) By `Matt Savoie `_, - `Justus Magin `_ and `Tom Nicholas - `_. -- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By - `Matt Savoie `_ and `Tom Nicholas +- Imports ``datatree`` repository and history into internal location. (:pull:`8688`) + By `Matt Savoie `_, `Justus Magin `_ + and `Tom Nicholas `_. +- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) + By `Matt Savoie `_ and `Tom Nicholas `_. - Refactor :py:meth:`xarray.core.indexing.DaskIndexingAdapter.__getitem__` to remove an unnecessary rewrite of the indexer key (:issue: `8377`, :pull:`8758`) diff --git a/xarray/tests/datatree/test_treenode.py b/xarray/tests/test_treenode.py similarity index 100% rename from xarray/tests/datatree/test_treenode.py rename to xarray/tests/test_treenode.py From 4830cd4b3170a753e0b877c86c7a6a9b7aaf3f23 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Tue, 20 Feb 2024 16:27:18 -0700 Subject: [PATCH 13/19] refactors backend tests for datatree IO --- xarray/tests/conftest.py | 62 ++++++++++++++++++++++++ xarray/tests/datatree/conftest.py | 65 -------------------------- xarray/tests/{datatree => }/test_io.py | 65 +++++++++++--------------- 3 files changed, 89 insertions(+), 103 deletions(-) delete mode 100644 xarray/tests/datatree/conftest.py rename xarray/tests/{datatree => }/test_io.py (70%) diff --git a/xarray/tests/conftest.py b/xarray/tests/conftest.py index 7c30759e499..8590c9fb4e7 100644 --- a/xarray/tests/conftest.py +++ b/xarray/tests/conftest.py @@ -6,6 +6,7 @@ import xarray as xr from xarray import DataArray, Dataset +from xarray.datatree_.datatree import DataTree from xarray.tests import create_test_data, requires_dask @@ -136,3 +137,64 @@ def d(request, backend, type) -> DataArray | Dataset: return result else: raise ValueError + + +@pytest.fixture(scope="module") +def create_test_datatree(): + """ + Create a test datatree with this structure: + + + |-- set1 + | |-- + | | Dimensions: () + | | Data variables: + | | a int64 0 + | | b int64 1 + | |-- set1 + | |-- set2 + |-- set2 + | |-- + | | Dimensions: (x: 2) + | | Data variables: + | | a (x) int64 2, 3 + | | b (x) int64 0.1, 0.2 + | |-- set1 + |-- set3 + |-- + | Dimensions: (x: 2, y: 3) + | Data variables: + | a (y) int64 6, 7, 8 + | set0 (x) int64 9, 10 + + The structure has deliberately repeated names of tags, variables, and + dimensions in order to better check for bugs caused by name conflicts. + """ + + def _create_test_datatree(modify=lambda ds: ds): + set1_data = modify(xr.Dataset({"a": 0, "b": 1})) + set2_data = modify(xr.Dataset({"a": ("x", [2, 3]), "b": ("x", [0.1, 0.2])})) + root_data = modify(xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])})) + + # Avoid using __init__ so we can independently test it + root: DataTree = DataTree(data=root_data) + set1: DataTree = DataTree(name="set1", parent=root, data=set1_data) + DataTree(name="set1", parent=set1) + DataTree(name="set2", parent=set1) + set2: DataTree = DataTree(name="set2", parent=root, data=set2_data) + DataTree(name="set1", parent=set2) + DataTree(name="set3", parent=root) + + return root + + return _create_test_datatree + + +@pytest.fixture(scope="module") +def simple_datatree(create_test_datatree): + """ + Invoke create_test_datatree fixture (callback). + + Returns a DataTree. + """ + return create_test_datatree() diff --git a/xarray/tests/datatree/conftest.py b/xarray/tests/datatree/conftest.py deleted file mode 100644 index b341f3007aa..00000000000 --- a/xarray/tests/datatree/conftest.py +++ /dev/null @@ -1,65 +0,0 @@ -import pytest - -import xarray as xr -from xarray.datatree_.datatree import DataTree - - -@pytest.fixture(scope="module") -def create_test_datatree(): - """ - Create a test datatree with this structure: - - - |-- set1 - | |-- - | | Dimensions: () - | | Data variables: - | | a int64 0 - | | b int64 1 - | |-- set1 - | |-- set2 - |-- set2 - | |-- - | | Dimensions: (x: 2) - | | Data variables: - | | a (x) int64 2, 3 - | | b (x) int64 0.1, 0.2 - | |-- set1 - |-- set3 - |-- - | Dimensions: (x: 2, y: 3) - | Data variables: - | a (y) int64 6, 7, 8 - | set0 (x) int64 9, 10 - - The structure has deliberately repeated names of tags, variables, and - dimensions in order to better check for bugs caused by name conflicts. - """ - - def _create_test_datatree(modify=lambda ds: ds): - set1_data = modify(xr.Dataset({"a": 0, "b": 1})) - set2_data = modify(xr.Dataset({"a": ("x", [2, 3]), "b": ("x", [0.1, 0.2])})) - root_data = modify(xr.Dataset({"a": ("y", [6, 7, 8]), "set0": ("x", [9, 10])})) - - # Avoid using __init__ so we can independently test it - root: DataTree = DataTree(data=root_data) - set1: DataTree = DataTree(name="set1", parent=root, data=set1_data) - DataTree(name="set1", parent=set1) - DataTree(name="set2", parent=set1) - set2: DataTree = DataTree(name="set2", parent=root, data=set2_data) - DataTree(name="set1", parent=set2) - DataTree(name="set3", parent=root) - - return root - - return _create_test_datatree - - -@pytest.fixture(scope="module") -def simple_datatree(create_test_datatree): - """ - Invoke create_test_datatree fixture (callback). - - Returns a DataTree. - """ - return create_test_datatree() diff --git a/xarray/tests/datatree/test_io.py b/xarray/tests/test_io.py similarity index 70% rename from xarray/tests/datatree/test_io.py rename to xarray/tests/test_io.py index 4f32e19de4a..5291f88111d 100644 --- a/xarray/tests/datatree/test_io.py +++ b/xarray/tests/test_io.py @@ -9,68 +9,62 @@ ) -class TestIO: - @requires_netCDF4 +class DatatreeIOBase: + engine = None + def test_to_netcdf(self, tmpdir, simple_datatree): - filepath = str( - tmpdir / "test.nc" - ) # casting to str avoids a pathlib bug in xarray + filepath = tmpdir / "test.nc" original_dt = simple_datatree - original_dt.to_netcdf(filepath, engine="netcdf4") + original_dt.to_netcdf(filepath, engine=self.engine) - roundtrip_dt = open_datatree(filepath) + roundtrip_dt = open_datatree(filepath, engine=self.engine) assert_equal(original_dt, roundtrip_dt) - @requires_netCDF4 def test_netcdf_encoding(self, tmpdir, simple_datatree): - filepath = str( - tmpdir / "test.nc" - ) # casting to str avoids a pathlib bug in xarray + filepath = tmpdir / "test.nc" original_dt = simple_datatree # add compression comp = dict(zlib=True, complevel=9) enc = {"/set2": {var: comp for var in original_dt["/set2"].ds.data_vars}} - original_dt.to_netcdf(filepath, encoding=enc, engine="netcdf4") - roundtrip_dt = open_datatree(filepath) + original_dt.to_netcdf(filepath, encoding=enc, engine=self.engine) + roundtrip_dt = open_datatree(filepath, engine=self.engine) assert roundtrip_dt["/set2/a"].encoding["zlib"] == comp["zlib"] assert roundtrip_dt["/set2/a"].encoding["complevel"] == comp["complevel"] enc["/not/a/group"] = {"foo": "bar"} # type: ignore with pytest.raises(ValueError, match="unexpected encoding group.*"): - original_dt.to_netcdf(filepath, encoding=enc, engine="netcdf4") + original_dt.to_netcdf(filepath, encoding=enc, engine=self.engine) - @requires_h5netcdf - def test_to_h5netcdf(self, tmpdir, simple_datatree): - filepath = str( - tmpdir / "test.nc" - ) # casting to str avoids a pathlib bug in xarray - original_dt = simple_datatree - original_dt.to_netcdf(filepath, engine="h5netcdf") - roundtrip_dt = open_datatree(filepath) - assert_equal(original_dt, roundtrip_dt) +@requires_netCDF4 +class TestNetCDF4DatatreeIO(DatatreeIOBase): + engine = "netcdf4" + + +@requires_h5netcdf +class TestH5NetCDFDatatreeIO(DatatreeIOBase): + engine = "h5netcdf" + + +@requires_zarr +class TestZarrDatatreeIO: + engine = "zarr" - @requires_zarr def test_to_zarr(self, tmpdir, simple_datatree): - filepath = str( - tmpdir / "test.zarr" - ) # casting to str avoids a pathlib bug in xarray + filepath = tmpdir / "test.zarr" original_dt = simple_datatree original_dt.to_zarr(filepath) - roundtrip_dt = open_datatree(filepath, engine="zarr") + roundtrip_dt = open_datatree(filepath) assert_equal(original_dt, roundtrip_dt) - @requires_zarr def test_zarr_encoding(self, tmpdir, simple_datatree): import zarr - filepath = str( - tmpdir / "test.zarr" - ) # casting to str avoids a pathlib bug in xarray + filepath = tmpdir / "test.zarr" original_dt = simple_datatree comp = {"compressor": zarr.Blosc(cname="zstd", clevel=3, shuffle=2)} @@ -85,13 +79,10 @@ def test_zarr_encoding(self, tmpdir, simple_datatree): with pytest.raises(ValueError, match="unexpected encoding group.*"): original_dt.to_zarr(filepath, encoding=enc, engine="zarr") - @requires_zarr def test_to_zarr_zip_store(self, tmpdir, simple_datatree): from zarr.storage import ZipStore - filepath = str( - tmpdir / "test.zarr.zip" - ) # casting to str avoids a pathlib bug in xarray + filepath = tmpdir / "test.zarr.zip" original_dt = simple_datatree store = ZipStore(filepath) original_dt.to_zarr(store) @@ -99,7 +90,6 @@ def test_to_zarr_zip_store(self, tmpdir, simple_datatree): roundtrip_dt = open_datatree(store, engine="zarr") assert_equal(original_dt, roundtrip_dt) - @requires_zarr def test_to_zarr_not_consolidated(self, tmpdir, simple_datatree): filepath = tmpdir / "test.zarr" zmetadata = filepath / ".zmetadata" @@ -114,7 +104,6 @@ def test_to_zarr_not_consolidated(self, tmpdir, simple_datatree): roundtrip_dt = open_datatree(filepath, engine="zarr") assert_equal(original_dt, roundtrip_dt) - @requires_zarr def test_to_zarr_default_write_mode(self, tmpdir, simple_datatree): import zarr From e8459bc45e720faf4f687444918cc414703baacb Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Tue, 20 Feb 2024 16:30:15 -0700 Subject: [PATCH 14/19] Add explicit engine back in test_to_zarr --- xarray/tests/test_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_io.py b/xarray/tests/test_io.py index 5291f88111d..ee4b90f0f8a 100644 --- a/xarray/tests/test_io.py +++ b/xarray/tests/test_io.py @@ -58,7 +58,7 @@ def test_to_zarr(self, tmpdir, simple_datatree): original_dt = simple_datatree original_dt.to_zarr(filepath) - roundtrip_dt = open_datatree(filepath) + roundtrip_dt = open_datatree(filepath, engine="zarr") assert_equal(original_dt, roundtrip_dt) def test_zarr_encoding(self, tmpdir, simple_datatree): From 4238ce297d7bfef133c44efc037f8212cea24288 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 21 Feb 2024 10:00:32 -0700 Subject: [PATCH 15/19] Removes OrderedDict from treenode --- xarray/core/treenode.py | 42 +++++++++++++++++++---------------------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index c408c43d740..e401ff36c2c 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -1,7 +1,6 @@ from __future__ import annotations import sys -from collections import OrderedDict from collections.abc import Iterator, Mapping from pathlib import PurePosixPath from typing import ( @@ -51,8 +50,8 @@ class TreeNode(Generic[Tree]): This class stores no data, it has only parents and children attributes, and various methods. - Stores child nodes in an Ordered Dictionary, which is necessary to ensure that equality checks between two trees - also check that the order of child nodes is the same. + Stores child nodes in an dict, ensuring that equality checks between trees + and order of child nodes is the preserved (since python 3.7). Nodes themselves are intrinsically unnamed (do not possess a ._name attribute), but if the node has a parent you can find the key it is stored under via the .name property. @@ -69,15 +68,16 @@ class TreeNode(Generic[Tree]): Also allows access to any other node in the tree via unix-like paths, including upwards referencing via '../'. (This class is heavily inspired by the anytree library's NodeMixin class.) + """ _parent: Tree | None - _children: OrderedDict[str, Tree] + _children: dict[str, Tree] def __init__(self, children: Mapping[str, Tree] | None = None): """Create a parentless node.""" self._parent = None - self._children = OrderedDict() + self._children = {} if children is not None: self.children = children @@ -123,13 +123,11 @@ def _detach(self, parent: Tree | None) -> None: if parent is not None: self._pre_detach(parent) parents_children = parent.children - parent._children = OrderedDict( - { - name: child - for name, child in parents_children.items() - if child is not self - } - ) + parent._children = { + name: child + for name, child in parents_children.items() + if child is not self + } self._parent = None self._post_detach(parent) @@ -163,7 +161,7 @@ def children(self: Tree) -> Mapping[str, Tree]: @children.setter def children(self: Tree, children: Mapping[str, Tree]) -> None: self._check_children(children) - children = OrderedDict(children) + children = {**children} old_children = self.children del self.children @@ -311,20 +309,18 @@ def leaves(self: Tree) -> tuple[Tree, ...]: return tuple([node for node in self.subtree if node.is_leaf]) @property - def siblings(self: Tree) -> OrderedDict[str, Tree]: + def siblings(self: Tree) -> dict[str, Tree]: """ Nodes with the same parent as this node. """ if self.parent: - return OrderedDict( - { - name: child - for name, child in self.parent.children.items() - if child is not self - } - ) + return { + name: child + for name, child in self.parent.children.items() + if child is not self + } else: - return OrderedDict() + return {} @property def subtree(self: Tree) -> Iterator[Tree]: @@ -580,7 +576,7 @@ class NamedNode(TreeNode, Generic[Tree]): _name: str | None _parent: Tree | None - _children: OrderedDict[str, Tree] + _children: dict[str, Tree] def __init__(self, name=None, children=None): super().__init__(children=children) From db264c0453c6a1658315a0a9e78cfbf011f3699f Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 21 Feb 2024 10:00:55 -0700 Subject: [PATCH 16/19] Renames tests/test_io.py -> tests/test_backends_datatree.py --- xarray/tests/{test_io.py => test_backends_datatree.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename xarray/tests/{test_io.py => test_backends_datatree.py} (100%) diff --git a/xarray/tests/test_io.py b/xarray/tests/test_backends_datatree.py similarity index 100% rename from xarray/tests/test_io.py rename to xarray/tests/test_backends_datatree.py From b4acb0d6a559c42d367fd2e4329f399a8c1fc4de Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 21 Feb 2024 10:18:00 -0700 Subject: [PATCH 17/19] typo --- xarray/core/treenode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index e401ff36c2c..b3e6e43f306 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -51,7 +51,7 @@ class TreeNode(Generic[Tree]): This class stores no data, it has only parents and children attributes, and various methods. Stores child nodes in an dict, ensuring that equality checks between trees - and order of child nodes is the preserved (since python 3.7). + and order of child nodes is preserved (since python 3.7). Nodes themselves are intrinsically unnamed (do not possess a ._name attribute), but if the node has a parent you can find the key it is stored under via the .name property. From 0f4b38ae48ddf1635d4eb3e6fa2fbf3168c5e6bb Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 21 Feb 2024 11:00:20 -0700 Subject: [PATCH 18/19] Add types --- xarray/tests/test_backends_datatree.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py index ee4b90f0f8a..41a5297ca1a 100644 --- a/xarray/tests/test_backends_datatree.py +++ b/xarray/tests/test_backends_datatree.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + import pytest from xarray.backends.api import open_datatree @@ -8,9 +10,12 @@ requires_zarr, ) +if TYPE_CHECKING: + from xarray.backends.api import T_NetcdfEngine + class DatatreeIOBase: - engine = None + engine: T_NetcdfEngine | None = None def test_to_netcdf(self, tmpdir, simple_datatree): filepath = tmpdir / "test.nc" @@ -41,12 +46,12 @@ def test_netcdf_encoding(self, tmpdir, simple_datatree): @requires_netCDF4 class TestNetCDF4DatatreeIO(DatatreeIOBase): - engine = "netcdf4" + engine: T_NetcdfEngine | None = "netcdf4" @requires_h5netcdf class TestH5NetCDFDatatreeIO(DatatreeIOBase): - engine = "h5netcdf" + engine: T_NetcdfEngine | None = "h5netcdf" @requires_zarr From f2f327f7e65e9634215e6841e96210176803b013 Mon Sep 17 00:00:00 2001 From: Matt Savoie Date: Wed, 21 Feb 2024 13:22:31 -0700 Subject: [PATCH 19/19] Pass mypy for 3.9 --- xarray/tests/test_backends_datatree.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py index 41a5297ca1a..7bdb2b532d9 100644 --- a/xarray/tests/test_backends_datatree.py +++ b/xarray/tests/test_backends_datatree.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import TYPE_CHECKING import pytest