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]