File size: 6,498 Bytes
4cef980 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
"""
Text-based visual representations of graphs
"""
__all__ = ["forest_str"]
def forest_str(graph, with_labels=True, sources=None, write=None, ascii_only=False):
"""
Creates a nice utf8 representation of a directed forest
Parameters
----------
graph : nx.DiGraph | nx.Graph
Graph to represent (must be a tree, forest, or the empty graph)
with_labels : bool
If True will use the "label" attribute of a node to display if it
exists otherwise it will use the node value itself. Defaults to True.
sources : List
Mainly relevant for undirected forests, specifies which nodes to list
first. If unspecified the root nodes of each tree will be used for
directed forests; for undirected forests this defaults to the nodes
with the smallest degree.
write : callable
Function to use to write to, if None new lines are appended to
a list and returned. If set to the `print` function, lines will
be written to stdout as they are generated. If specified,
this function will return None. Defaults to None.
ascii_only : Boolean
If True only ASCII characters are used to construct the visualization
Returns
-------
str | None :
utf8 representation of the tree / forest
Example
-------
>>> graph = nx.balanced_tree(r=2, h=3, create_using=nx.DiGraph)
>>> print(nx.forest_str(graph))
╙── 0
├─╼ 1
│ ├─╼ 3
│ │ ├─╼ 7
│ │ └─╼ 8
│ └─╼ 4
│ ├─╼ 9
│ └─╼ 10
└─╼ 2
├─╼ 5
│ ├─╼ 11
│ └─╼ 12
└─╼ 6
├─╼ 13
└─╼ 14
>>> graph = nx.balanced_tree(r=1, h=2, create_using=nx.Graph)
>>> print(nx.forest_str(graph))
╙── 0
└── 1
└── 2
>>> print(nx.forest_str(graph, ascii_only=True))
+-- 0
L-- 1
L-- 2
"""
import networkx as nx
printbuf = []
if write is None:
_write = printbuf.append
else:
_write = write
# Define glphys
# Notes on available box and arrow characters
# https://en.wikipedia.org/wiki/Box-drawing_character
# https://stackoverflow.com/questions/2701192/triangle-arrow
if ascii_only:
glyph_empty = "+"
glyph_newtree_last = "+-- "
glyph_newtree_mid = "+-- "
glyph_endof_forest = " "
glyph_within_forest = ": "
glyph_within_tree = "| "
glyph_directed_last = "L-> "
glyph_directed_mid = "|-> "
glyph_undirected_last = "L-- "
glyph_undirected_mid = "|-- "
else:
glyph_empty = "╙"
glyph_newtree_last = "╙── "
glyph_newtree_mid = "╟── "
glyph_endof_forest = " "
glyph_within_forest = "╎ "
glyph_within_tree = "│ "
glyph_directed_last = "└─╼ "
glyph_directed_mid = "├─╼ "
glyph_undirected_last = "└── "
glyph_undirected_mid = "├── "
if len(graph.nodes) == 0:
_write(glyph_empty)
else:
if not nx.is_forest(graph):
raise nx.NetworkXNotImplemented("input must be a forest or the empty graph")
is_directed = graph.is_directed()
succ = graph.succ if is_directed else graph.adj
if sources is None:
if is_directed:
# use real source nodes for directed trees
sources = [n for n in graph.nodes if graph.in_degree[n] == 0]
else:
# use arbitrary sources for undirected trees
sources = [
min(cc, key=lambda n: graph.degree[n])
for cc in nx.connected_components(graph)
]
# Populate the stack with each source node, empty indentation, and mark
# the final node. Reverse the stack so sources are popped in the
# correct order.
last_idx = len(sources) - 1
stack = [(node, "", (idx == last_idx)) for idx, node in enumerate(sources)][
::-1
]
seen = set()
while stack:
node, indent, islast = stack.pop()
if node in seen:
continue
seen.add(node)
if not indent:
# Top level items (i.e. trees in the forest) get different
# glyphs to indicate they are not actually connected
if islast:
this_prefix = indent + glyph_newtree_last
next_prefix = indent + glyph_endof_forest
else:
this_prefix = indent + glyph_newtree_mid
next_prefix = indent + glyph_within_forest
else:
# For individual tree edges distinguish between directed and
# undirected cases
if is_directed:
if islast:
this_prefix = indent + glyph_directed_last
next_prefix = indent + glyph_endof_forest
else:
this_prefix = indent + glyph_directed_mid
next_prefix = indent + glyph_within_tree
else:
if islast:
this_prefix = indent + glyph_undirected_last
next_prefix = indent + glyph_endof_forest
else:
this_prefix = indent + glyph_undirected_mid
next_prefix = indent + glyph_within_tree
if with_labels:
label = graph.nodes[node].get("label", node)
else:
label = node
_write(this_prefix + str(label))
# Push children on the stack in reverse order so they are popped in
# the original order.
children = [child for child in succ[node] if child not in seen]
for idx, child in enumerate(children[::-1], start=1):
islast_next = idx <= 1
try_frame = (child, next_prefix, islast_next)
stack.append(try_frame)
if write is None:
# Only return a string if the custom write function was not specified
return "\n".join(printbuf)
|