Visualizing a tree with nodes of different colors with anytree

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.

Thank you very much for your help!

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.sstatic.net/lhPqT.png