From 320a8a2872b327e6f8e15be673a7e06d8028b037 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 10 Sep 2024 22:20:48 -0400 Subject: [PATCH 01/11] sketch tests --- xarray/tests/test_treenode.py | 55 ++++++++++++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_treenode.py b/xarray/tests/test_treenode.py index d9d581cc314..9b0d2efa474 100644 --- a/xarray/tests/test_treenode.py +++ b/xarray/tests/test_treenode.py @@ -1,12 +1,19 @@ from __future__ import annotations from collections.abc import Iterator +from textwrap import dedent from typing import cast import pytest from xarray.core.iterators import LevelOrderIter -from xarray.core.treenode import InvalidTreeError, NamedNode, NodePath, TreeNode +from xarray.core.treenode import ( + Children, + InvalidTreeError, + NamedNode, + NodePath, + TreeNode, +) class TestFamilyTree: @@ -224,6 +231,52 @@ def test_overwrite_child(self): assert marys_evil_twin.parent is john +class TestChildren: + def test_properties(self): + sue: TreeNode = TreeNode() + kate: TreeNode = TreeNode() + mary = TreeNode(children={"Sue": sue, "Kate": kate}) + + children = mary.children + assert isinstance(children, Children) + + # len + assert len(children) == 2 + + # iter + assert list(children) == ["Sue", "Kate"] + + assert mary.children["Sue"] is sue + assert mary.children["Kate"] is kate + + assert "Sue" in mary.children + assert "Kate" in mary.children + assert 0 not in mary.children + assert "foo" not in mary.children + + with pytest.raises(KeyError): + children["foo"] + with pytest.raises(KeyError): + children[0] + + # TODO not sure what this should look like... + # repr + expected = dedent( + """\ + Children: + * x (x) int64 16B -1 -2 + * y (y) int64 24B 0 1 2 + a (x) int64 16B 4 5 + b int64 8B -10""" + ) + # actual = repr(coords) + # assert expected == actual + + def test_modify(self): + # TODO + ... + + class TestPruning: def test_del_child(self): john: TreeNode = TreeNode() From 1383211289a5c2d2043441271e7e8b6e8905b4e3 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 10 Sep 2024 22:20:56 -0400 Subject: [PATCH 02/11] sketch implementation --- xarray/core/treenode.py | 69 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 66 insertions(+), 3 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 84ce392ad32..2b7f545ff14 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -5,11 +5,12 @@ from pathlib import PurePosixPath from typing import ( TYPE_CHECKING, + Any, Generic, TypeVar, ) -from xarray.core.utils import Frozen, is_dict_like +from xarray.core.utils import is_dict_like if TYPE_CHECKING: from xarray.core.types import T_DataArray @@ -44,6 +45,68 @@ def __init__(self, *pathsegments): Tree = TypeVar("Tree", bound="TreeNode") +class Children(Mapping[str, Tree], Generic[Tree]): + """ + Dictionary-like container for the immediate children of a single DataTree node. + + This collection can be passed directly to the :py:class:`~xarray.DataTree` constructor via its `children` argument. + """ + + _treenode: Tree + + # TODO add slots? + # __slots__ = ("_data",) + + def __init__(self, treenode: Tree): + self._treenode = treenode + + def _names(self) -> set[str]: + return set(self._treenode.children.keys()) + + def __iter__(self) -> Iterator[str]: + return iter(self._names) + + def __len__(self) -> int: + return len(self._names) + + def __contains__(self, key: str) -> bool: + return key in self._names + + def __repr__(self) -> str: + # TODO + from xarray.core import formatting + + return formatting.children_repr(self) + + def __getitem__(self, key: str) -> Tree: + return self._treenode._children[key] + + def __delitem__(self, key: str) -> None: + if key in self._names: + child = self._treenode._children[key] + del self._treenode._children[key] + child.orphan() + else: + raise KeyError(key) + + def __setitem__(self, key: str, value: Any) -> None: + self.update({key: value}) + + def update(self, other: Mapping[Any, Any]) -> None: + """Update with other child nodes.""" + # TODO + ... + + def _ipython_key_completions_(self): + """Provide method for the key-autocompletions in IPython.""" + # TODO + return [ + key + for key in self._treenode._ipython_key_completions_() + if key not in self._treenode.variables + ] + + class TreeNode(Generic[Tree]): """ Base class representing a node of a tree, with methods for traversing and altering the tree. @@ -160,9 +223,9 @@ def orphan(self) -> None: self._set_parent(new_parent=None) @property - def children(self: Tree) -> Mapping[str, Tree]: + def children(self: Tree) -> Children[str, Tree]: """Child nodes of this node, stored under a mapping via their names.""" - return Frozen(self._children) + return Children(self) @children.setter def children(self: Tree, children: Mapping[str, Tree]) -> None: From e61112f900166d501bed0af1f6d8b8ee9bf46736 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 10 Sep 2024 22:32:53 -0400 Subject: [PATCH 03/11] make ._names a property --- xarray/core/treenode.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 2b7f545ff14..9b2084ce5d8 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -60,6 +60,7 @@ class Children(Mapping[str, Tree], Generic[Tree]): def __init__(self, treenode: Tree): self._treenode = treenode + @property def _names(self) -> set[str]: return set(self._treenode.children.keys()) From 2c1b13f8d18ac2c0cd78af7648fbf6803edd95a8 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 10 Sep 2024 22:33:57 -0400 Subject: [PATCH 04/11] call private ._children --- 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 9b2084ce5d8..c05c4b42b63 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -62,7 +62,7 @@ def __init__(self, treenode: Tree): @property def _names(self) -> set[str]: - return set(self._treenode.children.keys()) + return set(self._treenode._children.keys()) def __iter__(self) -> Iterator[str]: return iter(self._names) From 03c56ea197e713e70899cc63ee4850df79af28ca Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Tue, 10 Sep 2024 22:42:32 -0400 Subject: [PATCH 05/11] draft implementation of .update --- xarray/core/treenode.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index c05c4b42b63..14b089ec140 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -93,10 +93,17 @@ def __delitem__(self, key: str) -> None: def __setitem__(self, key: str, value: Any) -> None: self.update({key: value}) - def update(self, other: Mapping[Any, Any]) -> None: + def update(self, other: Mapping[str, Tree]) -> None: """Update with other child nodes.""" - # TODO - ... + + # TODO forbid strings with `/` in here? + + if not len(other): + return + + children = self._treenode._children.copy() + children.update(other) + self._treenode.children = children def _ipython_key_completions_(self): """Provide method for the key-autocompletions in IPython.""" From 860764d2c578476a9b47f5633386f06978725aff Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 08:16:41 -0400 Subject: [PATCH 06/11] fix bug with ordering of child names --- xarray/core/treenode.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 14b089ec140..7df46fe36e8 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -61,8 +61,8 @@ def __init__(self, treenode: Tree): self._treenode = treenode @property - def _names(self) -> set[str]: - return set(self._treenode._children.keys()) + def _names(self) -> list[str]: + return list(self._treenode._children.keys()) def __iter__(self) -> Iterator[str]: return iter(self._names) From 51c198ea360c78f4288cf6975aeb700d689ffe86 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 08:17:03 -0400 Subject: [PATCH 07/11] completely comment out repr test for now --- xarray/tests/test_treenode.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/xarray/tests/test_treenode.py b/xarray/tests/test_treenode.py index 9b0d2efa474..f289af3e277 100644 --- a/xarray/tests/test_treenode.py +++ b/xarray/tests/test_treenode.py @@ -1,7 +1,6 @@ from __future__ import annotations from collections.abc import Iterator -from textwrap import dedent from typing import cast import pytest @@ -261,14 +260,14 @@ def test_properties(self): # TODO not sure what this should look like... # repr - expected = dedent( - """\ - Children: - * x (x) int64 16B -1 -2 - * y (y) int64 24B 0 1 2 - a (x) int64 16B 4 5 - b int64 8B -10""" - ) + # expected = dedent( + # """\ + # Children: + # * x (x) int64 16B -1 -2 + # * y (y) int64 24B 0 1 2 + # a (x) int64 16B 4 5 + # b int64 8B -10""" + # ) # actual = repr(coords) # assert expected == actual From e0358aa0b5cf08d231ea856d871700c1a9dbcfb7 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 08:22:40 -0400 Subject: [PATCH 08/11] check only immediate children are accessible --- xarray/tests/test_treenode.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/xarray/tests/test_treenode.py b/xarray/tests/test_treenode.py index f289af3e277..4b7651c968b 100644 --- a/xarray/tests/test_treenode.py +++ b/xarray/tests/test_treenode.py @@ -233,25 +233,29 @@ def test_overwrite_child(self): class TestChildren: def test_properties(self): sue: TreeNode = TreeNode() + mary: TreeNode = TreeNode(children={"Sue": sue}) kate: TreeNode = TreeNode() - mary = TreeNode(children={"Sue": sue, "Kate": kate}) + john = TreeNode(children={"Mary": mary, "Kate": kate}) - children = mary.children + children = john.children assert isinstance(children, Children) # len assert len(children) == 2 # iter - assert list(children) == ["Sue", "Kate"] + assert list(children) == ["Mary", "Kate"] - assert mary.children["Sue"] is sue - assert mary.children["Kate"] is kate + assert john.children["Mary"] is mary + assert john.children["Kate"] is kate + + assert "Mary" in john.children + assert "Kate" in john.children + assert 0 not in john.children + assert "foo" not in john.children - assert "Sue" in mary.children - assert "Kate" in mary.children - assert 0 not in mary.children - assert "foo" not in mary.children + # only immediate children should be accessible + assert "sue" not in john.children with pytest.raises(KeyError): children["foo"] From 08f83f8d1aaca856ed9f6febebd71ccf5fb8f888 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 09:21:13 -0400 Subject: [PATCH 09/11] add tests for modifying via .children --- xarray/tests/test_treenode.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_treenode.py b/xarray/tests/test_treenode.py index 4b7651c968b..921108a5f88 100644 --- a/xarray/tests/test_treenode.py +++ b/xarray/tests/test_treenode.py @@ -276,7 +276,28 @@ def test_properties(self): # assert expected == actual def test_modify(self): - # TODO + sue: TreeNode = TreeNode() + mary: TreeNode = TreeNode(children={"Sue": sue}) + kate: TreeNode = TreeNode() + john = TreeNode(children={"Mary": mary, "Kate": kate}) + + children = john.children + + # test assignment + ashley: TreeNode = TreeNode() + children["Ashley"] = ashley + assert john.children["Ashley"] is ashley + + # test deletion + del children["Ashley"] + assert "Ashley" not in john.children + + # test constructor + john2 = TreeNode(children=children) + assert john2.children == children + + def test_modify_below_root(self): + # TODO test that modifying .children doesn't affect grandparent ... From 7f87d0846c0510fdb1e4510e5d73d34efa19afd6 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 09:21:25 -0400 Subject: [PATCH 10/11] temporary repr for debugging --- xarray/core/treenode.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 7df46fe36e8..49c5921e633 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -74,10 +74,13 @@ def __contains__(self, key: str) -> bool: return key in self._names def __repr__(self) -> str: - # TODO - from xarray.core import formatting + return f"Children({dict(self._treenode._children)})" + + # def __repr__(self) -> str: + # # TODO + # from xarray.core import formatting - return formatting.children_repr(self) + # return formatting.children_repr(self) def __getitem__(self, key: str) -> Tree: return self._treenode._children[key] From d6de4a18e6df5f3ad32cf0095b1dd8b8e1699b73 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Wed, 11 Sep 2024 17:27:22 -0400 Subject: [PATCH 11/11] implement a very simple repr --- xarray/core/treenode.py | 8 +------- xarray/tests/test_treenode.py | 20 +++++++++----------- 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 49c5921e633..932cab2fa23 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -74,13 +74,7 @@ def __contains__(self, key: str) -> bool: return key in self._names def __repr__(self) -> str: - return f"Children({dict(self._treenode._children)})" - - # def __repr__(self) -> str: - # # TODO - # from xarray.core import formatting - - # return formatting.children_repr(self) + return "\n".join(["Children:"] + [f" {name}" for name in self._names]) def __getitem__(self, key: str) -> Tree: return self._treenode._children[key] diff --git a/xarray/tests/test_treenode.py b/xarray/tests/test_treenode.py index 921108a5f88..02d478c24de 100644 --- a/xarray/tests/test_treenode.py +++ b/xarray/tests/test_treenode.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator +from textwrap import dedent from typing import cast import pytest @@ -262,18 +263,15 @@ def test_properties(self): with pytest.raises(KeyError): children[0] - # TODO not sure what this should look like... # repr - # expected = dedent( - # """\ - # Children: - # * x (x) int64 16B -1 -2 - # * y (y) int64 24B 0 1 2 - # a (x) int64 16B 4 5 - # b int64 8B -10""" - # ) - # actual = repr(coords) - # assert expected == actual + expected = dedent( + """\ + Children: + Mary + Kate""" + ) + actual = repr(children) + assert expected == actual def test_modify(self): sue: TreeNode = TreeNode()