0

I am looking for a simple way to construct and draw a tree (on google colab).

Importantly, I would like to have nodes of different colors and shapes. Ideally, I would like something as follows.

from anytree import Node, RenderTree
from anytree.exporter import DotExporter
from IPython.display import Image


# construct tree
ceo = Node("CEO") #root

vp_1 = Node("VP_1", parent=ceo, color="red")
vp_2 = Node("VP_2", parent=ceo)

gm_1 = Node("GM_1", parent=vp_1, shape="square", color="red")
gm_2 = Node("GM_2", parent=vp_2, shape="square")

m_1 = Node("M_1", parent=gm_2)

# draw tree
DotExporter(ceo).to_picture("ceo.png")

# show image
Image('ceo.png')

As color and shape are not real arguments of Node, this code currently generates the following image. I would like VP_1 and GM_1 to be red, and GM_1 and GM_2 to be squares.

enter image description here

Thank you very much for your help!

grg
  • 85
  • 7
  • anytree doesn't support reshaping and recoloring from what i see in the docs – rikyeah Jan 17 '22 at 21:13
  • i'm ok with using something other than anytree. – grg Jan 18 '22 at 10:34
  • 1
    then https://stackoverflow.com/questions/2442014/tree-libraries-in-python might be a good starting point to search libraries that support different visualization formats. – rikyeah Jan 18 '22 at 17:20
  • 1
    Thanks! I ended up using graphviz: https://graphviz.readthedocs.io/en/stable/examples.html – grg Jan 26 '22 at 16:28

1 Answers1

4

Anytree's DotExporter has a nodeattrfunc argument, where you can pass a function that accepts a Node and returns the attributes it should be given in a DOT language string. Thus, if you have stored your color and shape information as node attributes, you can use a custom nodeattrfunc like the following to translate those into DOT attributes:

def set_color_shape(node):
    attrs = []
    attrs += [f'color={node.color}'] if hasattr(node, 'color') else []
    attrs += [f'shape={node.shape}'] if hasattr(node, 'shape') else []
    return ', '.join(attrs)

DotExporter(ceo, nodeattrfunc=set_color_shape).to_picture('example.png')

For the result, see: https://i.stack.imgur.com/lhPqT.png

txp
  • 66
  • 2