def test_map_children(): """Test recursive application of map on children""" # a # | \ \ \ # b c d e # | \ | # f g h # # Should become # # 2 # | \ \ \ # 3 3 3 3 # | \ | # 4 4 4 root = TreeNode(1) b = TreeNode(1, root) TreeNode(1, root) TreeNode(1, root) e = TreeNode(1, root) TreeNode(1, b) TreeNode(1, b) TreeNode(1, e) def increment(node, children): for child in TreeNode(0, None, children): child.item += 1 return node.item + 1, children rval = root.map(increment, root.children) assert [node.item for node in rval] == [2, 3, 4, 4, 3, 3, 3, 4]
def test_map_parent(): """Test recursive application of map on parents""" # a # | \ \ \ # b c d e # | \ | # f g h # # Will give # # 4 # | # 3 # | # 2 root = TreeNode(1) b = TreeNode(1, root) TreeNode(1, root) TreeNode(1, root) e = TreeNode(1, root) f = TreeNode(1, b) TreeNode(1, b) h = TreeNode(1, e) def increment_parent(node, parent): if parent is not None: for parent in parent.root: parent.item += 1 return node.item + 1, parent rval = f.map(increment_parent, f.parent) assert [node.item for node in rval.root] == [4, 3, 2] rval = h.map(increment_parent, h.parent) assert [node.item for node in rval.root] == [4, 3, 2]