513 lines
15 KiB
Python
513 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
import pytest
|
|
|
|
from xarray.core.treenode import (
|
|
InvalidTreeError,
|
|
NamedNode,
|
|
NodePath,
|
|
TreeNode,
|
|
group_subtrees,
|
|
zip_subtrees,
|
|
)
|
|
|
|
|
|
class TestFamilyTree:
|
|
def test_lonely(self) -> None:
|
|
root: TreeNode = TreeNode()
|
|
assert root.parent is None
|
|
assert root.children == {}
|
|
|
|
def test_parenting(self) -> None:
|
|
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) -> None:
|
|
john: TreeNode = TreeNode()
|
|
|
|
with pytest.raises(InvalidTreeError, match="cannot be a parent of itself"):
|
|
john._set_parent(john, "John")
|
|
|
|
with pytest.raises(InvalidTreeError, match="cannot be a parent of itself"):
|
|
john.children = {"John": john}
|
|
|
|
mary: TreeNode = TreeNode()
|
|
rose: TreeNode = TreeNode()
|
|
mary._set_parent(john, "Mary")
|
|
rose._set_parent(mary, "Rose")
|
|
|
|
with pytest.raises(InvalidTreeError, match="is already a descendant"):
|
|
john._set_parent(rose, "John")
|
|
|
|
with pytest.raises(InvalidTreeError, match="is already a descendant"):
|
|
rose.children = {"John": john}
|
|
|
|
def test_parent_swap(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
mary._set_parent(john, "Mary")
|
|
|
|
steve: TreeNode = TreeNode()
|
|
mary._set_parent(steve, "Mary")
|
|
|
|
assert mary.parent == steve
|
|
assert steve.children["Mary"] is mary
|
|
assert "Mary" not in john.children
|
|
|
|
def test_forbid_setting_parent_directly(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
|
|
with pytest.raises(
|
|
AttributeError, match="Cannot set parent attribute directly"
|
|
):
|
|
mary.parent = john
|
|
|
|
def test_dont_modify_children_inplace(self) -> None:
|
|
# GH issue 9196
|
|
child: TreeNode = TreeNode()
|
|
TreeNode(children={"child": child})
|
|
assert child.parent is None
|
|
|
|
def test_multi_child_family(self) -> None:
|
|
john: TreeNode = TreeNode(children={"Mary": TreeNode(), "Kate": TreeNode()})
|
|
|
|
assert "Mary" in john.children
|
|
mary = john.children["Mary"]
|
|
assert isinstance(mary, TreeNode)
|
|
assert mary.parent is john
|
|
|
|
assert "Kate" in john.children
|
|
kate = john.children["Kate"]
|
|
assert isinstance(kate, TreeNode)
|
|
assert kate.parent is john
|
|
|
|
def test_disown_child(self) -> None:
|
|
john: TreeNode = TreeNode(children={"Mary": TreeNode()})
|
|
mary = john.children["Mary"]
|
|
mary.orphan()
|
|
assert mary.parent is None
|
|
assert "Mary" not in john.children
|
|
|
|
def test_doppelganger_child(self) -> None:
|
|
kate: TreeNode = TreeNode()
|
|
john: TreeNode = TreeNode()
|
|
|
|
with pytest.raises(TypeError):
|
|
john.children = {"Kate": 666}
|
|
|
|
with pytest.raises(InvalidTreeError, match="Cannot add same node"):
|
|
john.children = {"Kate": kate, "Evil_Kate": kate}
|
|
|
|
john = TreeNode(children={"Kate": kate})
|
|
evil_kate: TreeNode = TreeNode()
|
|
evil_kate._set_parent(john, "Kate")
|
|
assert john.children["Kate"] is evil_kate
|
|
|
|
def test_sibling_relationships(self) -> None:
|
|
john: TreeNode = TreeNode(
|
|
children={"Mary": TreeNode(), "Kate": TreeNode(), "Ashley": TreeNode()}
|
|
)
|
|
kate = john.children["Kate"]
|
|
assert list(kate.siblings) == ["Mary", "Ashley"]
|
|
assert "Kate" not in kate.siblings
|
|
|
|
def test_copy_subtree(self) -> None:
|
|
tony: TreeNode = TreeNode()
|
|
michael: TreeNode = TreeNode(children={"Tony": tony})
|
|
vito = TreeNode(children={"Michael": michael})
|
|
|
|
# check that children of assigned children are also copied (i.e. that ._copy_subtree works)
|
|
copied_tony = vito.children["Michael"].children["Tony"]
|
|
assert copied_tony is not tony
|
|
|
|
def test_parents(self) -> None:
|
|
vito: TreeNode = TreeNode(
|
|
children={"Michael": TreeNode(children={"Tony": TreeNode()})},
|
|
)
|
|
michael = vito.children["Michael"]
|
|
tony = michael.children["Tony"]
|
|
|
|
assert tony.root is vito
|
|
assert tony.parents == (michael, vito)
|
|
|
|
|
|
class TestGetNodes:
|
|
def test_get_child(self) -> None:
|
|
john: TreeNode = TreeNode(
|
|
children={
|
|
"Mary": TreeNode(
|
|
children={"Sue": TreeNode(children={"Steven": TreeNode()})}
|
|
)
|
|
}
|
|
)
|
|
mary = john.children["Mary"]
|
|
sue = mary.children["Sue"]
|
|
steven = sue.children["Steven"]
|
|
|
|
# get child
|
|
assert john._get_item("Mary") is mary
|
|
assert mary._get_item("Sue") is sue
|
|
|
|
# no child exists
|
|
with pytest.raises(KeyError):
|
|
john._get_item("Kate")
|
|
|
|
# get grandchild
|
|
assert john._get_item("Mary/Sue") is sue
|
|
|
|
# get great-grandchild
|
|
assert john._get_item("Mary/Sue/Steven") is steven
|
|
|
|
# get from middle of tree
|
|
assert mary._get_item("Sue/Steven") is steven
|
|
|
|
def test_get_upwards(self) -> None:
|
|
john: TreeNode = TreeNode(
|
|
children={
|
|
"Mary": TreeNode(children={"Sue": TreeNode(), "Kate": TreeNode()})
|
|
}
|
|
)
|
|
mary = john.children["Mary"]
|
|
sue = mary.children["Sue"]
|
|
kate = mary.children["Kate"]
|
|
|
|
assert sue._get_item("../") is mary
|
|
assert sue._get_item("../../") is john
|
|
|
|
# relative path
|
|
assert sue._get_item("../Kate") is kate
|
|
|
|
def test_get_from_root(self) -> None:
|
|
john: TreeNode = TreeNode(
|
|
children={"Mary": TreeNode(children={"Sue": TreeNode()})}
|
|
)
|
|
mary = john.children["Mary"]
|
|
sue = mary.children["Sue"]
|
|
|
|
assert sue._get_item("/Mary") is mary
|
|
|
|
|
|
class TestSetNodes:
|
|
def test_set_child_node(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
john._set_item("Mary", mary)
|
|
|
|
assert john.children["Mary"] is mary
|
|
assert isinstance(mary, TreeNode)
|
|
assert mary.children == {}
|
|
assert mary.parent is john
|
|
|
|
def test_child_already_exists(self) -> None:
|
|
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) -> None:
|
|
rose: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
john: TreeNode = TreeNode()
|
|
|
|
john._set_item("Mary", mary)
|
|
john._set_item("Mary/Rose", rose)
|
|
|
|
assert john.children["Mary"] is mary
|
|
assert isinstance(mary, TreeNode)
|
|
assert "Rose" in mary.children
|
|
assert rose.parent is mary
|
|
|
|
def test_create_intermediate_child(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
rose: TreeNode = TreeNode()
|
|
|
|
# test intermediate children not allowed
|
|
with pytest.raises(KeyError, match="Could not reach"):
|
|
john._set_item(path="Mary/Rose", item=rose, new_nodes_along_path=False)
|
|
|
|
# test intermediate children allowed
|
|
john._set_item("Mary/Rose", rose, new_nodes_along_path=True)
|
|
assert "Mary" in john.children
|
|
mary = john.children["Mary"]
|
|
assert isinstance(mary, TreeNode)
|
|
assert mary.children == {"Rose": rose}
|
|
assert rose.parent == mary
|
|
assert rose.parent == mary
|
|
|
|
def test_overwrite_child(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
john._set_item("Mary", mary)
|
|
|
|
# test overwriting not allowed
|
|
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
|
|
assert marys_evil_twin.parent is None
|
|
|
|
# test overwriting allowed
|
|
marys_evil_twin = TreeNode()
|
|
john._set_item("Mary", marys_evil_twin, allow_overwrite=True)
|
|
assert john.children["Mary"] is marys_evil_twin
|
|
assert marys_evil_twin.parent is john
|
|
|
|
|
|
class TestPruning:
|
|
def test_del_child(self) -> None:
|
|
john: TreeNode = TreeNode()
|
|
mary: TreeNode = TreeNode()
|
|
john._set_item("Mary", mary)
|
|
|
|
del john["Mary"]
|
|
assert "Mary" not in john.children
|
|
assert mary.parent is None
|
|
|
|
with pytest.raises(KeyError):
|
|
del john["Mary"]
|
|
|
|
|
|
def create_test_tree() -> tuple[NamedNode, NamedNode]:
|
|
# a
|
|
# ├── b
|
|
# │ ├── d
|
|
# │ └── e
|
|
# │ ├── f
|
|
# │ └── g
|
|
# └── c
|
|
# └── h
|
|
# └── i
|
|
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}
|
|
e.children = {"f": f, "g": g}
|
|
c.children = {"h": h}
|
|
h.children = {"i": i}
|
|
|
|
return a, f
|
|
|
|
|
|
class TestGroupSubtrees:
|
|
def test_one_tree(self) -> None:
|
|
root, _ = create_test_tree()
|
|
expected_names = [
|
|
"a",
|
|
"b",
|
|
"c",
|
|
"d",
|
|
"e",
|
|
"h",
|
|
"f",
|
|
"g",
|
|
"i",
|
|
]
|
|
expected_paths = [
|
|
".",
|
|
"b",
|
|
"c",
|
|
"b/d",
|
|
"b/e",
|
|
"c/h",
|
|
"b/e/f",
|
|
"b/e/g",
|
|
"c/h/i",
|
|
]
|
|
result_paths, result_names = zip(
|
|
*[(path, node.name) for path, (node,) in group_subtrees(root)], strict=False
|
|
)
|
|
assert list(result_names) == expected_names
|
|
assert list(result_paths) == expected_paths
|
|
|
|
result_names_ = [node.name for (node,) in zip_subtrees(root)]
|
|
assert result_names_ == expected_names
|
|
|
|
def test_different_order(self) -> None:
|
|
first: NamedNode = NamedNode(
|
|
name="a", children={"b": NamedNode(), "c": NamedNode()}
|
|
)
|
|
second: NamedNode = NamedNode(
|
|
name="a", children={"c": NamedNode(), "b": NamedNode()}
|
|
)
|
|
assert [node.name for node in first.subtree] == ["a", "b", "c"]
|
|
assert [node.name for node in second.subtree] == ["a", "c", "b"]
|
|
assert [(x.name, y.name) for x, y in zip_subtrees(first, second)] == [
|
|
("a", "a"),
|
|
("b", "b"),
|
|
("c", "c"),
|
|
]
|
|
assert [path for path, _ in group_subtrees(first, second)] == [".", "b", "c"]
|
|
|
|
def test_different_structure(self) -> None:
|
|
first: NamedNode = NamedNode(name="a", children={"b": NamedNode()})
|
|
second: NamedNode = NamedNode(name="a", children={"c": NamedNode()})
|
|
it = group_subtrees(first, second)
|
|
|
|
path, (node1, node2) = next(it)
|
|
assert path == "."
|
|
assert node1.name == node2.name == "a"
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=re.escape(r"children at root node do not match: ['b'] vs ['c']"),
|
|
):
|
|
next(it)
|
|
|
|
|
|
class TestAncestry:
|
|
def test_parents(self) -> None:
|
|
_, leaf_f = create_test_tree()
|
|
expected = ["e", "b", "a"]
|
|
assert [node.name for node in leaf_f.parents] == expected
|
|
|
|
def test_lineage(self) -> None:
|
|
_, leaf_f = create_test_tree()
|
|
expected = ["f", "e", "b", "a"]
|
|
with pytest.warns(DeprecationWarning):
|
|
assert [node.name for node in leaf_f.lineage] == expected
|
|
|
|
def test_ancestors(self) -> None:
|
|
_, leaf_f = create_test_tree()
|
|
with pytest.warns(DeprecationWarning):
|
|
ancestors = leaf_f.ancestors
|
|
expected = ["a", "b", "e", "f"]
|
|
for node, expected_name in zip(ancestors, expected, strict=True):
|
|
assert node.name == expected_name
|
|
|
|
def test_subtree(self) -> None:
|
|
root, _ = create_test_tree()
|
|
expected = [
|
|
"a",
|
|
"b",
|
|
"c",
|
|
"d",
|
|
"e",
|
|
"h",
|
|
"f",
|
|
"g",
|
|
"i",
|
|
]
|
|
actual = [node.name for node in root.subtree]
|
|
assert expected == actual
|
|
|
|
def test_subtree_with_keys(self) -> None:
|
|
root, _ = create_test_tree()
|
|
expected_names = [
|
|
"a",
|
|
"b",
|
|
"c",
|
|
"d",
|
|
"e",
|
|
"h",
|
|
"f",
|
|
"g",
|
|
"i",
|
|
]
|
|
expected_paths = [
|
|
".",
|
|
"b",
|
|
"c",
|
|
"b/d",
|
|
"b/e",
|
|
"c/h",
|
|
"b/e/f",
|
|
"b/e/g",
|
|
"c/h/i",
|
|
]
|
|
result_paths, result_names = zip(
|
|
*[(path, node.name) for path, node in root.subtree_with_keys], strict=False
|
|
)
|
|
assert list(result_names) == expected_names
|
|
assert list(result_paths) == expected_paths
|
|
|
|
def test_descendants(self) -> None:
|
|
root, _ = create_test_tree()
|
|
descendants = root.descendants
|
|
expected = [
|
|
"b",
|
|
"c",
|
|
"d",
|
|
"e",
|
|
"h",
|
|
"f",
|
|
"g",
|
|
"i",
|
|
]
|
|
for node, expected_name in zip(descendants, expected, strict=True):
|
|
assert node.name == expected_name
|
|
|
|
def test_leaves(self) -> None:
|
|
tree, _ = create_test_tree()
|
|
leaves = tree.leaves
|
|
expected = [
|
|
"d",
|
|
"f",
|
|
"g",
|
|
"i",
|
|
]
|
|
for node, expected_name in zip(leaves, expected, strict=True):
|
|
assert node.name == expected_name
|
|
|
|
def test_levels(self) -> None:
|
|
a, f = create_test_tree()
|
|
|
|
assert a.level == 0
|
|
assert f.level == 3
|
|
|
|
assert a.depth == 3
|
|
assert f.depth == 3
|
|
|
|
assert a.width == 1
|
|
assert f.width == 3
|
|
|
|
|
|
class TestRenderTree:
|
|
def test_render_nodetree(self) -> None:
|
|
john: NamedNode = NamedNode(
|
|
children={
|
|
"Mary": NamedNode(children={"Sam": NamedNode(), "Ben": NamedNode()}),
|
|
"Kate": NamedNode(),
|
|
}
|
|
)
|
|
mary = john.children["Mary"]
|
|
|
|
expected_nodes = [
|
|
"NamedNode()",
|
|
"\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_nodes = john_repr.splitlines()
|
|
assert len(john_nodes) == len(expected_nodes)
|
|
for expected_node, repr_node in zip(expected_nodes, john_nodes, strict=True):
|
|
assert expected_node == repr_node
|
|
|
|
|
|
def test_nodepath():
|
|
path = NodePath("/Mary")
|
|
assert path.root == "/"
|
|
assert path.stem == "Mary"
|