#!/usr/bin/env python # -*- coding: utf-8 -*- from __future__ import unicode_literals import sys import os import unittest from treelib import Tree, Node from treelib.tree import NodeIDAbsentError, LoopError def encode(value): if sys.version_info[0] == 2: # Python2.x : return value.encode("utf-8") else: # Python3.x : return value class TreeCase(unittest.TestCase): def setUp(self): tree = Tree(identifier="tree 1") tree.create_node("Hárry", "hárry") tree.create_node("Jane", "jane", parent="hárry") tree.create_node("Bill", "bill", parent="hárry") tree.create_node("Diane", "diane", parent="jane") tree.create_node("George", "george", parent="bill") # Hárry # |-- Jane # |-- Diane # |-- Bill # |-- George self.tree = tree self.copytree = Tree(self.tree, deep=True) self.input_dict = { "Bill": "Harry", "Jane": "Harry", "Harry": None, "Diane": "Jane", "Mark": "Jane", "Mary": "Harry", } @staticmethod def get_t1(): """ root ├── A │ └── A1 └── B """ t = Tree(identifier="t1") t.create_node(tag="root", identifier="r") t.create_node(tag="A", identifier="a", parent="r") t.create_node(tag="B", identifier="b", parent="r") t.create_node(tag="A1", identifier="a1", parent="a") return t @staticmethod def get_t2(): """ root2 ├── C └── D └── D1 """ t = Tree(identifier="t2") t.create_node(tag="root2", identifier="r2") t.create_node(tag="C", identifier="c", parent="r2") t.create_node(tag="D", identifier="d", parent="r2") t.create_node(tag="D1", identifier="d1", parent="d") return t def test_tree(self): self.assertEqual(isinstance(self.tree, Tree), True) self.assertEqual(isinstance(self.copytree, Tree), True) def test_is_root(self): # retro-compatibility self.assertTrue(self.tree._nodes["hárry"].is_root()) self.assertFalse(self.tree._nodes["jane"].is_root()) def test_tree_wise_is_root(self): subtree = self.tree.subtree("jane", identifier="subtree 2") # harry is root of tree 1 but not present in subtree 2 self.assertTrue(self.tree._nodes["hárry"].is_root("tree 1")) self.assertNotIn("hárry", subtree._nodes) # jane is not root of tree 1 but is root of subtree 2 self.assertFalse(self.tree._nodes["jane"].is_root("tree 1")) self.assertTrue(subtree._nodes["jane"].is_root("subtree 2")) def test_paths_to_leaves(self): paths = self.tree.paths_to_leaves() self.assertEqual(len(paths), 2) self.assertTrue(["hárry", "jane", "diane"] in paths) self.assertTrue(["hárry", "bill", "george"] in paths) def test_nodes(self): self.assertEqual(len(self.tree.nodes), 5) self.assertEqual(len(self.tree.all_nodes()), 5) self.assertEqual(self.tree.size(), 5) self.assertEqual(self.tree.get_node("jane").tag, "Jane") self.assertEqual(self.tree.contains("jane"), True) self.assertEqual("jane" in self.tree, True) self.assertEqual(self.tree.contains("alien"), False) self.tree.create_node("Alien", "alien", parent="jane") self.assertEqual(self.tree.contains("alien"), True) self.tree.remove_node("alien") def test_getitem(self): """Nodes can be accessed via getitem.""" for node_id in self.tree.nodes: try: self.tree[node_id] except NodeIDAbsentError: self.fail("Node access should be possible via getitem.") try: self.tree["root"] except NodeIDAbsentError: pass else: self.fail("There should be no default fallback value for getitem.") def test_parent(self): for nid in self.tree.nodes: if nid == self.tree.root: self.assertEqual(self.tree.parent(nid), None) else: self.assertEqual(self.tree.parent(nid) in self.tree.all_nodes(), True) def test_ancestor(self): for nid in self.tree.nodes: if nid == self.tree.root: self.assertEqual(self.tree.ancestor(nid), None) else: for level in range(self.tree.level(nid) - 1, 0, -1): self.assertEqual( self.tree.ancestor(nid, level=level) in self.tree.all_nodes(), True, ) def test_children(self): for nid in self.tree.nodes: children = self.tree.is_branch(nid) for child in children: self.assertEqual(self.tree[child] in self.tree.all_nodes(), True) children = self.tree.children(nid) for child in children: self.assertEqual(child in self.tree.all_nodes(), True) try: self.tree.is_branch("alien") except NodeIDAbsentError: pass else: self.fail("The absent node should be declaimed.") def test_remove_node(self): self.tree.create_node("Jill", "jill", parent="george") self.tree.create_node("Mark", "mark", parent="jill") self.assertEqual(self.tree.remove_node("jill"), 2) self.assertEqual(self.tree.get_node("jill") is None, True) self.assertEqual(self.tree.get_node("mark") is None, True) def test_tree_wise_depth(self): # Try getting the level of this tree self.assertEqual(self.tree.depth(), 2) self.tree.create_node("Jill", "jill", parent="george") self.assertEqual(self.tree.depth(), 3) self.tree.create_node("Mark", "mark", parent="jill") self.assertEqual(self.tree.depth(), 4) # Try getting the level of the node """ self.tree.show() Hárry |___ Bill | |___ George | |___ Jill | |___ Mark |___ Jane | |___ Diane """ self.assertEqual(self.tree.depth(self.tree.get_node("mark")), 4) self.assertEqual(self.tree.depth(self.tree.get_node("jill")), 3) self.assertEqual(self.tree.depth(self.tree.get_node("george")), 2) self.assertEqual(self.tree.depth("jane"), 1) self.assertEqual(self.tree.depth("bill"), 1) self.assertEqual(self.tree.depth("hárry"), 0) # Try getting Exception node = Node("Test One", "identifier 1") self.assertRaises(NodeIDAbsentError, self.tree.depth, node) # Reset the test case self.tree.remove_node("jill") def test_leaves(self): # retro-compatibility leaves = self.tree.leaves() for nid in self.tree.expand_tree(): self.assertEqual( (self.tree[nid].is_leaf()) == (self.tree[nid] in leaves), True ) leaves = self.tree.leaves(nid="jane") for nid in self.tree.expand_tree(nid="jane"): self.assertEqual( self.tree[nid].is_leaf() == (self.tree[nid] in leaves), True ) def test_tree_wise_leaves(self): leaves = self.tree.leaves() for nid in self.tree.expand_tree(): self.assertEqual( (self.tree[nid].is_leaf("tree 1")) == (self.tree[nid] in leaves), True ) leaves = self.tree.leaves(nid="jane") for nid in self.tree.expand_tree(nid="jane"): self.assertEqual( self.tree[nid].is_leaf("tree 1") == (self.tree[nid] in leaves), True ) def test_link_past_node(self): self.tree.create_node("Jill", "jill", parent="hárry") self.tree.create_node("Mark", "mark", parent="jill") self.assertEqual("mark" not in self.tree.is_branch("hárry"), True) self.tree.link_past_node("jill") self.assertEqual("mark" in self.tree.is_branch("hárry"), True) def test_expand_tree(self): # default config # Hárry # |-- Jane # |-- Diane # |-- Bill # |-- George # Traverse in depth first mode preserving insertion order nodes = [nid for nid in self.tree.expand_tree(sorting=False)] self.assertEqual(nodes, ["h\xe1rry", "jane", "diane", "bill", "george"]) self.assertEqual(len(nodes), 5) # By default traverse depth first and sort child nodes by node tag nodes = [nid for nid in self.tree.expand_tree()] self.assertEqual(nodes, ["h\xe1rry", "bill", "george", "jane", "diane"]) self.assertEqual(len(nodes), 5) # expanding from specific node nodes = [nid for nid in self.tree.expand_tree(nid="bill")] self.assertEqual(nodes, ["bill", "george"]) self.assertEqual(len(nodes), 2) # changing into width mode preserving insertion order nodes = [nid for nid in self.tree.expand_tree(mode=Tree.WIDTH, sorting=False)] self.assertEqual(nodes, ["h\xe1rry", "jane", "bill", "diane", "george"]) self.assertEqual(len(nodes), 5) # Breadth first mode, child nodes sorting by tag nodes = [nid for nid in self.tree.expand_tree(mode=Tree.WIDTH)] self.assertEqual(nodes, ["h\xe1rry", "bill", "jane", "george", "diane"]) self.assertEqual(len(nodes), 5) # expanding by filters # Stops at root nodes = [nid for nid in self.tree.expand_tree(filter=lambda x: x.tag == "Bill")] self.assertEqual(len(nodes), 0) nodes = [nid for nid in self.tree.expand_tree(filter=lambda x: x.tag != "Bill")] self.assertEqual(nodes, ["h\xe1rry", "jane", "diane"]) self.assertEqual(len(nodes), 3) def test_move_node(self): diane_parent = self.tree.parent("diane") self.tree.move_node("diane", "bill") self.assertEqual("diane" in self.tree.is_branch("bill"), True) self.tree.move_node("diane", diane_parent.identifier) def test_paste_tree(self): new_tree = Tree() new_tree.create_node("Jill", "jill") new_tree.create_node("Mark", "mark", parent="jill") self.tree.paste("jane", new_tree) self.assertEqual("jill" in self.tree.is_branch("jane"), True) self.tree.show() self.assertEqual( self.tree._reader, """Hárry ├── Bill │ └── George └── Jane ├── Diane └── Jill └── Mark """, ) self.tree.remove_node("jill") self.assertNotIn("jill", self.tree.nodes.keys()) self.assertNotIn("mark", self.tree.nodes.keys()) self.tree.show() self.assertEqual( self.tree._reader, """Hárry ├── Bill │ └── George └── Jane └── Diane """, ) def test_merge(self): # merge on empty initial tree t1 = Tree(identifier="t1") t2 = self.get_t2() t1.merge(nid=None, new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r2") self.assertEqual(set(t1._nodes.keys()), {"r2", "c", "d", "d1"}) self.assertEqual( t1.show(stdout=False), """root2 ├── C └── D └── D1 """, ) # merge empty new_tree (on root) t1 = self.get_t1() t2 = Tree(identifier="t2") t1.merge(nid="r", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertEqual(set(t1._nodes.keys()), {"r", "a", "a1", "b"}) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 └── B """, ) # merge at root t1 = self.get_t1() t2 = self.get_t2() t1.merge(nid="r", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertNotIn("r2", t1._nodes.keys()) self.assertEqual(set(t1._nodes.keys()), {"r", "a", "a1", "b", "c", "d", "d1"}) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 ├── B ├── C └── D └── D1 """, ) # merge on node t1 = self.get_t1() t2 = self.get_t2() t1.merge(nid="b", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertNotIn("r2", t1._nodes.keys()) self.assertEqual(set(t1._nodes.keys()), {"r", "a", "a1", "b", "c", "d", "d1"}) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 └── B ├── C └── D └── D1 """, ) def test_paste(self): # paste under root t1 = self.get_t1() t2 = self.get_t2() t1.paste(nid="r", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertEqual(t1.parent("r2").identifier, "r") self.assertEqual( set(t1._nodes.keys()), {"r", "r2", "a", "a1", "b", "c", "d", "d1"} ) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 ├── B └── root2 ├── C └── D └── D1 """, ) # paste under non-existing node t1 = self.get_t1() t2 = self.get_t2() with self.assertRaises(NodeIDAbsentError) as e: t1.paste(nid="not_existing", new_tree=t2) self.assertEqual(e.exception.args[0], "Node 'not_existing' is not in the tree") # paste under None nid t1 = self.get_t1() t2 = self.get_t2() with self.assertRaises(ValueError) as e: t1.paste(nid=None, new_tree=t2) self.assertEqual( e.exception.args[0], 'Must define "nid" under which new tree is pasted.' ) # paste under node t1 = self.get_t1() t2 = self.get_t2() t1.paste(nid="b", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertEqual(t1.parent("b").identifier, "r") self.assertEqual( set(t1._nodes.keys()), {"r", "a", "a1", "b", "c", "d", "d1", "r2"} ) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 └── B └── root2 ├── C └── D └── D1 """, ) # paste empty new_tree (under root) t1 = self.get_t1() t2 = Tree(identifier="t2") t1.paste(nid="r", new_tree=t2) self.assertEqual(t1.identifier, "t1") self.assertEqual(t1.root, "r") self.assertEqual(set(t1._nodes.keys()), {"r", "a", "a1", "b"}) self.assertEqual( t1.show(stdout=False), """root ├── A │ └── A1 └── B """, ) def test_rsearch(self): for nid in ["hárry", "jane", "diane"]: self.assertEqual(nid in self.tree.rsearch("diane"), True) def test_subtree(self): subtree_copy = Tree(self.tree.subtree("jane"), deep=True) self.assertEqual(subtree_copy.parent("jane") is None, True) subtree_copy["jane"].tag = "Sweeti" self.assertEqual(self.tree["jane"].tag == "Jane", True) self.assertEqual(subtree_copy.level("diane"), 1) self.assertEqual(subtree_copy.level("jane"), 0) self.assertEqual(self.tree.level("jane"), 1) def test_remove_subtree(self): subtree_shallow = self.tree.remove_subtree("jane") self.assertEqual("jane" not in self.tree.is_branch("hárry"), True) self.tree.paste("hárry", subtree_shallow) def test_remove_subtree_whole_tree(self): self.tree.remove_subtree("hárry") self.assertIsNone(self.tree.root) self.assertEqual(len(self.tree.nodes.keys()), 0) def test_to_json(self): self.assertEqual.__self__.maxDiff = None self.tree.to_json() self.tree.to_json(True) def test_siblings(self): self.assertEqual(len(self.tree.siblings("hárry")) == 0, True) self.assertEqual(self.tree.siblings("jane")[0].identifier == "bill", True) def test_tree_data(self): class Flower(object): def __init__(self, color): self.color = color self.tree.create_node("Jill", "jill", parent="jane", data=Flower("white")) self.assertEqual(self.tree["jill"].data.color, "white") self.tree.remove_node("jill") def test_show_data_property(self): new_tree = Tree() sys.stdout = open(os.devnull, "w") # stops from printing to console try: new_tree.show() class Flower(object): def __init__(self, color): self.color = color new_tree.create_node("Jill", "jill", data=Flower("white")) new_tree.show(data_property="color") finally: sys.stdout.close() sys.stdout = sys.__stdout__ # stops from printing to console def test_level(self): self.assertEqual(self.tree.level("hárry"), 0) depth = self.tree.depth() self.assertEqual(self.tree.level("diane"), depth) self.assertEqual( self.tree.level("diane", lambda x: x.identifier != "jane"), depth - 1 ) def test_size(self): self.assertEqual(self.tree.size(level=2), 2) self.assertEqual(self.tree.size(level=1), 2) self.assertEqual(self.tree.size(level=0), 1) def test_print_backend(self): expected_result = """\ Hárry ├── Bill │ └── George └── Jane └── Diane """ assert str(self.tree) == encode(expected_result) def test_show(self): if sys.version_info[0] < 3: reload(sys) # noqa: F821 sys.setdefaultencoding("utf-8") sys.stdout = open(os.devnull, "w") # stops from printing to console try: self.tree.show() finally: sys.stdout.close() sys.stdout = sys.__stdout__ # stops from printing to console def tearDown(self): self.tree = None self.copytree = None def test_show_without_sorting(self): t = Tree() t.create_node("Students", "Students", parent=None) Node(tag="Students", identifier="Students", data=None) t.create_node("Ben", "Ben", parent="Students") Node(tag="Ben", identifier="Ben", data=None) t.create_node("Annie", "Annie", parent="Students") Node(tag="Annie", identifier="Annie", data=None) t.show() self.assertEqual( t.show(sorting=False, stdout=False), """Students ├── Ben └── Annie """, ) def test_all_nodes_itr(self): """ tests: Tree.all_nodes_iter Added by: William Rusnack """ new_tree = Tree() self.assertEqual(len(new_tree.all_nodes_itr()), 0) nodes = list() nodes.append(new_tree.create_node("root_node")) nodes.append(new_tree.create_node("second", parent=new_tree.root)) for nd in new_tree.all_nodes_itr(): self.assertTrue(nd in nodes) def test_filter_nodes(self): """ tests: Tree.filter_nodes Added by: William Rusnack """ new_tree = Tree(identifier="tree 1") self.assertEqual(tuple(new_tree.filter_nodes(lambda n: True)), ()) nodes = list() nodes.append(new_tree.create_node("root_node")) nodes.append(new_tree.create_node("second", parent=new_tree.root)) self.assertEqual(tuple(new_tree.filter_nodes(lambda n: False)), ()) self.assertEqual( tuple(new_tree.filter_nodes(lambda n: n.is_root("tree 1"))), (nodes[0],) ) self.assertEqual( tuple(new_tree.filter_nodes(lambda n: not n.is_root("tree 1"))), (nodes[1],) ) self.assertTrue(set(new_tree.filter_nodes(lambda n: True)), set(nodes)) def test_loop(self): tree = Tree() tree.create_node("a", "a") tree.create_node("b", "b", parent="a") tree.create_node("c", "c", parent="b") tree.create_node("d", "d", parent="c") try: tree.move_node("b", "d") except LoopError: pass def test_modify_node_identifier_directly_failed(self): tree = Tree() tree.create_node("Harry", "harry") tree.create_node("Jane", "jane", parent="harry") n = tree.get_node("jane") self.assertTrue(n.identifier == "jane") # Failed to modify n.identifier = "xyz" self.assertTrue(tree.get_node("xyz") is None) self.assertTrue(tree.get_node("jane").identifier == "xyz") def test_modify_node_identifier_recursively(self): tree = Tree() tree.create_node("Harry", "harry") tree.create_node("Jane", "jane", parent="harry") n = tree.get_node("jane") self.assertTrue(n.identifier == "jane") # Success to modify tree.update_node(n.identifier, identifier="xyz") self.assertTrue(tree.get_node("jane") is None) self.assertTrue(tree.get_node("xyz").identifier == "xyz") def test_modify_node_identifier_root(self): tree = Tree(identifier="tree 3") tree.create_node("Harry", "harry") tree.create_node("Jane", "jane", parent="harry") tree.update_node(tree["harry"].identifier, identifier="xyz", tag="XYZ") self.assertTrue(tree.root == "xyz") self.assertTrue(tree["xyz"].tag == "XYZ") self.assertEqual(tree.parent("jane").identifier, "xyz") def test_subclassing(self): class SubNode(Node): pass class SubTree(Tree): node_class = SubNode tree = SubTree() node = tree.create_node() self.assertTrue(isinstance(node, SubNode)) tree = Tree(node_class=SubNode) node = tree.create_node() self.assertTrue(isinstance(node, SubNode)) def test_shallow_copy_hermetic_pointers(self): # tree 1 # Hárry # └── Jane # └── Diane # └── Bill # └── George tree2 = self.tree.subtree(nid="jane", identifier="tree 2") # tree 2 # Jane # └── Diane # check that in shallow copy, instances are the same self.assertIs(self.tree["jane"], tree2["jane"]) self.assertEqual( self.tree["jane"]._predecessor, {"tree 1": "hárry", "tree 2": None} ) self.assertEqual( dict(self.tree["jane"]._successors), {"tree 1": ["diane"], "tree 2": ["diane"]}, ) # when creating new node on subtree, check that it has no impact on initial tree tree2.create_node("Jill", "jill", parent="diane") self.assertIn("jill", tree2) self.assertIn("jill", tree2.is_branch("diane")) self.assertNotIn("jill", self.tree) self.assertNotIn("jill", self.tree.is_branch("diane")) def test_paste_duplicate_nodes(self): t1 = Tree() t1.create_node(identifier="A") t2 = Tree() t2.create_node(identifier="A") t2.create_node(identifier="B", parent="A") with self.assertRaises(ValueError) as e: t1.paste("A", t2) self.assertEqual(e.exception.args, ("Duplicated nodes ['A'] exists.",)) def test_shallow_paste(self): t1 = Tree() n1 = t1.create_node(identifier="A") t2 = Tree() n2 = t2.create_node(identifier="B") t3 = Tree() n3 = t3.create_node(identifier="C") t1.paste(n1.identifier, t2) self.assertEqual(t1.to_dict(), {"A": {"children": ["B"]}}) t1.paste(n1.identifier, t3) self.assertEqual(t1.to_dict(), {"A": {"children": ["B", "C"]}}) self.assertEqual(t1.level(n1.identifier), 0) self.assertEqual(t1.level(n2.identifier), 1) self.assertEqual(t1.level(n3.identifier), 1) def test_root_removal(self): t = Tree() t.create_node(identifier="root-A") self.assertEqual(len(t.nodes.keys()), 1) self.assertEqual(t.root, "root-A") t.remove_node(identifier="root-A") self.assertEqual(len(t.nodes.keys()), 0) self.assertEqual(t.root, None) t.create_node(identifier="root-B") self.assertEqual(len(t.nodes.keys()), 1) self.assertEqual(t.root, "root-B") def test_from_map(self): tree = Tree.from_map(self.input_dict) self.assertTrue(tree.size() == 6) self.assertTrue( tree.root == [k for k, v in self.input_dict.items() if v is None][0] ) tree = Tree.from_map(self.input_dict, id_func=lambda x: x.upper()) self.assertTrue(tree.size() == 6) def data_func(x): return x.upper() tree = Tree.from_map(self.input_dict, data_func=data_func) self.assertTrue(tree.size() == 6) self.assertTrue( tree.get_node(tree.root).data == data_func([k for k, v in self.input_dict.items() if v is None][0]) ) with self.assertRaises(ValueError): # invalid input payload without a root tree = Tree.from_map({"a": "b"}) with self.assertRaises(ValueError): # invalid input payload without more than 1 root tree = Tree.from_map({"a": None, "b": None})