Migrate treenode module.

doc/whats-new.rst

@@ -102,6 +102,8 @@ Internal Changes
 - Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By `Matt
   Savoie <https://github.com/flamingbear>`_.
 
+- Adds ``treenode`` functionality into ``xarray/core`` (:pull:`8757`) By `Matt Savoie <https://github.com/flamingbear>`_.
+
 .. _whats-new.2024.01.1:
 
 v2024.01.1 (23 Jan, 2024)

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():

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():

xarray/datatree_/datatree/treenode.py → 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
@@ -25,7 +21,7 @@ class InvalidTreeError(Exception):
 
 
 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):
@@ -75,10 +71,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 +87,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 +133,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 +238,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 +250,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 +262,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 +302,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 +337,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 +427,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 +441,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,24 +484,26 @@ 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:
         """
         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
         ------
@@ -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):

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__ = (

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

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."""
 

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):

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