Expose ability to individually style nodes in visualization

[skrawcz]

Context

From slack:

I want to be able to tag nodes with different types to give them a different color to show some more structure in the DAG.
Are there similar ideas or examples of this already?

What can be done today

You can post process the graph returned:

# main.py
import graphviz
from hamilton import driver

import functions

# customize graphviz render: https://graphviz.org/docs/nodes/
# careful with overwriting string attributes; fillcolor should be safe
level_stylesheet = dict(
    intermediate=dict(
        level="intermediate",  # add arbitrary metadata to the DOT file; could collide with graphviz attributes
        fillcolor="royalblue",  # edits the style
    ),
    final=dict(
        level="final",
        fillcolor="aquamarine",
    )
)


dr = driver.Builder().with_modules(functions).build()
g: graphviz.Digraph = dr.display_all_functions()

for v in dr.list_available_variables():
    if level := v.tags.get("level"):
        g.node(v.name, **level_stylesheet[level])

g

What could be done by the framework?

Idea 1: Expose callback

Expose a callback function to enable the above to be within the framework:

def apply_style(node_name: str, node_type: str, tags: dict) -> dict:
   """This function returns the graphviz styles to apply given the node name, node_type, and tags"""
   has_property = tags.get("custom_value")
   if has_property:
      return {"some": "values"}
   return {}

dr = driver.Builder().with_modules(functions).build()

dr.display_all_functions(node_style_callback=apply_style)

Idea 2: decorate functions

@custom_node_type(color=.., name=..)
def some_function() -> : 
     pass

Then when you use any of the display/visualize functions, it'll use what was annotated on the functions. But this
requires all functions that are to receive special treatment, to be annotated -- so it might mean that every function
has @custom_node_type?

Idea 3: do some combination of the above

Shift left is nice, but it could mean that things are verbose if applied to many functions. So we could allow/work with a combination. E.g. apply decorator first, and then apply any other functions.

Idea 4: Allow JSON export

We could also provide a more flexible export format for people to plug into their system of choice. "dot" files aren't that standard.
E.g. #635 show a prototype of this.

This option wouldn't conflict with having the others, as it would facilitate an escape hatch.

---

Got more ideas? add comments, and I'll update to expose them here. Else please provide comments.

Prototype

See current prototype here that does (3) - #642.

[skrawcz]

Questions:

• what about modifying the legend?

[zilto]

The following can be added to modify the legend. One limitation is that the "default styles" of the visualization aren't exposed in the internal API currently. They exists as nested functions of create_graphviz_graph. (ref)

the following can be added for the legend

# the style used for Function nodes
default_node_style = dict(
    shape="rectangle",
    margin="0.15",
    style="rounded,filled",
    fillcolor="#b4d8e4",
    fontname="Helvetica",
)

# `cluster__legend` is the name of the legend subgraph
with g.subgraph(name='cluster__legend') as legend_subgraph:
    for level, style in level_stylesheet.items():
        legend_node = dict(**default_node_style)  # set default style
        legend_node.update(**style)  # update default style with stylesheet
        legend_subgraph.node(level, label=level, **legend_node)

[roelbertens]

Would it be an option to create a decorator for this? Then you could also include a name to display in the legend for each node type. E.g. @custom_node_type(color=.., name=..)

Building upon the example from Thierry you would get something like this:

[skrawcz]

Yep I think that's a possibility. Having a decorator/tag annotation helps with the "shift left" theme of keeping everything together.

[skrawcz]

@roelbertens would this decorator be too much boilerplate if you need to annotate 100s of functions with it? 🤔

[roelbertens]

You have a point, but I'm doing the same with @tag now to define which nodes get which color.

[skrawcz]

Fair. We'll try to prototype something this week/weekend. I think we might want a combination.

[zilto]

Graphviz is a well supported and extensible renderer for graph objects. A key limitation is that many attributes, including style are defined using string with comma separated values (e.g., style="rounded,filled"). It becomes challenging to update styles and apply multiple conditional formatting because you need to parse the style string into a list of values to lookup each time.

Also, graphviz doesn't provide a way to read node data of a graphviz.Digraph object. To update a node, you need to completly overwrite it with a node of the same name. To remove it, you need to parse the strings of digraph_obj.body, which is a list of strings (each describing a node or an edge).

Anyhow, all of this to say that everything is simpler if we can style the entire graph "at once".

[skrawcz]

We could also use the tags to display things from the function: e.g. documentation & source code?

[alechewitt]

Another approach that could be beneficial is enabling the export of the graph into a JSON file.

Currently, users are limited to consuming the graph as an image file generated by Graphviz. This prevents the use of other visualization libraries or displaying the graph in text form.

One possible solution, is to expose a method on the driver that would create a JSON representation of the graph. It could work something like this:

# main.py
from hamilton import driver

import functions

dr = driver.Builder().with_modules(functions).build()

# Create JSON file of the Graph
dr.export_json("./my-graph.json")

Users could then use the JSON to create a visualization in anyway they choose. Eg: use D3 or display in an HTML table.

[skrawcz]

For other people following along, here's @alechewitt 's PR - #635

[skrawcz]

This has been shipped as part of 1.47.0 -- so I consider this discussion closed.