anusfoil's picture
Upload 3363 files
4cef980 verified
from textwrap import dedent
import pytest
import networkx as nx
def test_directed_tree_str():
# Create a directed forest with labels
graph = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
for node in graph.nodes:
graph.nodes[node]["label"] = "node_" + chr(ord("a") + node)
node_target = dedent(
"""
╙── 0
β”œβ”€β•Ό 1
β”‚Β Β  β”œβ”€β•Ό 3
β”‚Β Β  └─╼ 4
└─╼ 2
β”œβ”€β•Ό 5
└─╼ 6
"""
).strip()
label_target = dedent(
"""
╙── node_a
β”œβ”€β•Ό node_b
β”‚Β Β  β”œβ”€β•Ό node_d
β”‚Β Β  └─╼ node_e
└─╼ node_c
β”œβ”€β•Ό node_f
└─╼ node_g
"""
).strip()
# Basic node case
ret = nx.forest_str(graph, with_labels=False)
print(ret)
assert ret == node_target
# Basic label case
ret = nx.forest_str(graph, with_labels=True)
print(ret)
assert ret == label_target
# Custom write function case
lines = []
ret = nx.forest_str(graph, write=lines.append, with_labels=False)
assert ret is None
assert lines == node_target.split("\n")
# Smoke test to ensure passing the print function works. To properly test
# this case we would need to capture stdout. (for potential reference
# implementation see :class:`ubelt.util_stream.CaptureStdout`)
ret = nx.forest_str(graph, write=print)
assert ret is None
def test_empty_graph():
assert nx.forest_str(nx.DiGraph()) == "β•™"
assert nx.forest_str(nx.Graph()) == "β•™"
def test_directed_multi_tree_forest():
tree1 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
tree2 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
forest = nx.disjoint_union_all([tree1, tree2])
ret = nx.forest_str(forest)
print(ret)
target = dedent(
"""
β•Ÿβ”€β”€ 0
β•ŽΒ Β  β”œβ”€β•Ό 1
β•ŽΒ Β  β”‚Β Β  β”œβ”€β•Ό 3
β•ŽΒ Β  β”‚Β Β  └─╼ 4
β•ŽΒ Β  └─╼ 2
β•ŽΒ Β  β”œβ”€β•Ό 5
β•ŽΒ Β  └─╼ 6
╙── 7
β”œβ”€β•Ό 8
β”‚Β Β  β”œβ”€β•Ό 10
β”‚Β Β  └─╼ 11
└─╼ 9
β”œβ”€β•Ό 12
└─╼ 13
"""
).strip()
assert ret == target
tree3 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
forest = nx.disjoint_union_all([tree1, tree2, tree3])
ret = nx.forest_str(forest, sources=[0, 14, 7])
print(ret)
target = dedent(
"""
β•Ÿβ”€β”€ 0
β•ŽΒ Β  β”œβ”€β•Ό 1
β•ŽΒ Β  β”‚Β Β  β”œβ”€β•Ό 3
β•ŽΒ Β  β”‚Β Β  └─╼ 4
β•ŽΒ Β  └─╼ 2
β•ŽΒ Β  β”œβ”€β•Ό 5
β•ŽΒ Β  └─╼ 6
β•Ÿβ”€β”€ 14
β•ŽΒ Β  β”œβ”€β•Ό 15
β•ŽΒ Β  β”‚Β Β  β”œβ”€β•Ό 17
β•ŽΒ Β  β”‚Β Β  └─╼ 18
β•ŽΒ Β  └─╼ 16
β•ŽΒ Β  β”œβ”€β•Ό 19
β•ŽΒ Β  └─╼ 20
╙── 7
β”œβ”€β•Ό 8
β”‚Β Β  β”œβ”€β•Ό 10
β”‚Β Β  └─╼ 11
└─╼ 9
β”œβ”€β•Ό 12
└─╼ 13
"""
).strip()
assert ret == target
ret = nx.forest_str(forest, sources=[0, 14, 7], ascii_only=True)
print(ret)
target = dedent(
"""
+-- 0
:Β Β  |-> 1
:Β Β  |Β Β  |-> 3
:Β Β  |Β Β  L-> 4
:Β Β  L-> 2
:Β Β  |-> 5
:Β Β  L-> 6
+-- 14
:Β Β  |-> 15
:Β Β  |Β Β  |-> 17
:Β Β  |Β Β  L-> 18
:Β Β  L-> 16
:Β Β  |-> 19
:Β Β  L-> 20
+-- 7
|-> 8
|Β Β  |-> 10
|Β Β  L-> 11
L-> 9
|-> 12
L-> 13
"""
).strip()
assert ret == target
def test_undirected_multi_tree_forest():
tree1 = nx.balanced_tree(r=2, h=2, create_using=nx.Graph)
tree2 = nx.balanced_tree(r=2, h=2, create_using=nx.Graph)
tree2 = nx.relabel_nodes(tree2, {n: n + len(tree1) for n in tree2.nodes})
forest = nx.union(tree1, tree2)
ret = nx.forest_str(forest, sources=[0, 7])
print(ret)
target = dedent(
"""
β•Ÿβ”€β”€ 0
β•ŽΒ Β  β”œβ”€β”€ 1
β•ŽΒ Β  β”‚Β Β  β”œβ”€β”€ 3
β•ŽΒ Β  β”‚Β Β  └── 4
β•ŽΒ Β  └── 2
β•ŽΒ Β  β”œβ”€β”€ 5
β•ŽΒ Β  └── 6
╙── 7
β”œβ”€β”€ 8
β”‚Β Β  β”œβ”€β”€ 10
β”‚Β Β  └── 11
└── 9
β”œβ”€β”€ 12
└── 13
"""
).strip()
assert ret == target
ret = nx.forest_str(forest, sources=[0, 7], ascii_only=True)
print(ret)
target = dedent(
"""
+-- 0
:Β Β  |-- 1
:Β Β  |Β Β  |-- 3
:Β Β  |Β Β  L-- 4
:Β Β  L-- 2
:Β Β  |-- 5
:Β Β  L-- 6
+-- 7
|-- 8
|Β Β  |-- 10
|Β Β  L-- 11
L-- 9
|-- 12
L-- 13
"""
).strip()
assert ret == target
def test_undirected_tree_str():
# Create a directed forest with labels
graph = nx.balanced_tree(r=2, h=2, create_using=nx.Graph)
# arbitrary starting point
nx.forest_str(graph)
node_target0 = dedent(
"""
╙── 0
β”œβ”€β”€ 1
β”‚Β Β  β”œβ”€β”€ 3
β”‚Β Β  └── 4
└── 2
β”œβ”€β”€ 5
└── 6
"""
).strip()
# defined starting point
ret = nx.forest_str(graph, sources=[0])
print(ret)
assert ret == node_target0
# defined starting point
node_target2 = dedent(
"""
╙── 2
β”œβ”€β”€ 0
β”‚Β Β  └── 1
β”‚Β Β  β”œβ”€β”€ 3
β”‚Β Β  └── 4
β”œβ”€β”€ 5
└── 6
"""
).strip()
ret = nx.forest_str(graph, sources=[2])
print(ret)
assert ret == node_target2
def test_forest_str_errors():
ugraph = nx.complete_graph(3, create_using=nx.Graph)
with pytest.raises(nx.NetworkXNotImplemented):
nx.forest_str(ugraph)
dgraph = nx.complete_graph(3, create_using=nx.DiGraph)
with pytest.raises(nx.NetworkXNotImplemented):
nx.forest_str(dgraph)
def test_overspecified_sources():
"""
When sources are directly specified, we wont be able to determine when we
are in the last component, so there will always be a trailing, leftmost
pipe.
"""
graph = nx.disjoint_union_all(
[
nx.balanced_tree(r=2, h=1, create_using=nx.DiGraph),
nx.balanced_tree(r=1, h=2, create_using=nx.DiGraph),
nx.balanced_tree(r=2, h=1, create_using=nx.DiGraph),
]
)
# defined starting point
target1 = dedent(
"""
β•Ÿβ”€β”€ 0
β•ŽΒ Β  β”œβ”€β•Ό 1
β•ŽΒ Β  └─╼ 2
β•Ÿβ”€β”€ 3
β•ŽΒ Β  └─╼ 4
β•ŽΒ Β  └─╼ 5
β•Ÿβ”€β”€ 6
β•ŽΒ Β  β”œβ”€β•Ό 7
β•ŽΒ Β  └─╼ 8
"""
).strip()
target2 = dedent(
"""
β•Ÿβ”€β”€ 0
β•ŽΒ Β  β”œβ”€β•Ό 1
β•ŽΒ Β  └─╼ 2
β•Ÿβ”€β”€ 3
β•ŽΒ Β  └─╼ 4
β•ŽΒ Β  └─╼ 5
╙── 6
β”œβ”€β•Ό 7
└─╼ 8
"""
).strip()
lines = []
nx.forest_str(graph, write=lines.append, sources=graph.nodes)
got1 = chr(10).join(lines)
print("got1: ")
print(got1)
lines = []
nx.forest_str(graph, write=lines.append)
got2 = chr(10).join(lines)
print("got2: ")
print(got2)
assert got1 == target1
assert got2 == target2