From ee821b1cb12caec196eecf6e300935ae6aa89034 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Gyarmati?= <dev.petergy@gmail.com>
Date: Fri, 1 Mar 2024 13:56:08 +0100
Subject: [PATCH] feat(graphviz): support custom node and edge attributes in
 `ibis.visualize` (#8510)

---
 ibis/expr/types/core.py | 30 +++++++++++++++++++++++++++++-
 ibis/expr/visualize.py  | 27 ++++++++++++++++++++++++---
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/ibis/expr/types/core.py b/ibis/expr/types/core.py
index a665fbb397f7..f259e98d3258 100644
--- a/ibis/expr/types/core.py
+++ b/ibis/expr/types/core.py
@@ -163,6 +163,8 @@ def visualize(
         *,
         label_edges: bool = False,
         verbose: bool = False,
+        node_attr: Mapping[str, str] | None = None,
+        edge_attr: Mapping[str, str] | None = None,
     ) -> None:
         """Visualize an expression as a GraphViz graph in the browser.
 
@@ -175,6 +177,27 @@ def visualize(
             Show operation input names as edge labels
         verbose
             Print the graphviz DOT code to stderr if [](`True`)
+        node_attr
+            Mapping of ``(attribute, value)`` pairs set for all nodes.
+            Options are specified by the ``graphviz`` Python library.
+        edge_attr
+            Mapping of ``(attribute, value)`` pairs set for all edges.
+            Options are specified by the ``graphviz`` Python library.
+
+        Examples
+        --------
+        Open the visualization of an expression in default browser:
+
+        >>> import ibis
+        >>> left = ibis.table(dict(a="int64", b="string"), name="left")
+        >>> right = ibis.table(dict(b="string", c="int64", d="string"), name="right")
+        >>> expr = left.inner_join(right, "b").select(left.a, b=right.c, c=right.d)
+        >>> expr.visualize(
+        ...     format="svg",
+        ...     label_edges=True,
+        ...     node_attr={"fontname": "Roboto Mono", "fontsize": "10"},
+        ...     edge_attr={"fontsize": "8"},
+        ... )  # quartodoc: +SKIP # doctest: +SKIP
 
         Raises
         ------
@@ -184,7 +207,12 @@ def visualize(
         import ibis.expr.visualize as viz
 
         path = viz.draw(
-            viz.to_graph(self, label_edges=label_edges),
+            viz.to_graph(
+                self,
+                node_attr=node_attr,
+                edge_attr=edge_attr,
+                label_edges=label_edges,
+            ),
             format=format,
             verbose=verbose,
         )
diff --git a/ibis/expr/visualize.py b/ibis/expr/visualize.py
index ef16463251ce..bcaa41a53e38 100644
--- a/ibis/expr/visualize.py
+++ b/ibis/expr/visualize.py
@@ -93,8 +93,8 @@ def to_graph(expr, node_attr=None, edge_attr=None, label_edges: bool = False):
     graph = Graph.from_bfs(expr.op(), filter=ops.Node)
 
     g = gv.Digraph(
-        node_attr=node_attr or DEFAULT_NODE_ATTRS,
-        edge_attr=edge_attr or DEFAULT_EDGE_ATTRS,
+        node_attr=DEFAULT_NODE_ATTRS | (node_attr or {}),
+        edge_attr=DEFAULT_EDGE_ATTRS | (edge_attr or {}),
     )
 
     g.attr(rankdir="BT")
@@ -152,6 +152,7 @@ def draw(graph, path=None, format="png", verbose: bool = False):
 
 
 if __name__ == "__main__":
+    import json
     from argparse import ArgumentParser
 
     from ibis import _
@@ -173,6 +174,20 @@ def draw(graph, path=None, format="png", verbose: bool = False):
         action="store_true",
         help="Show operation inputs as edge labels.",
     )
+    p.add_argument(
+        "-n",
+        "--node-attr",
+        type=lambda x: json.loads(x) if x else {},
+        default="{}",
+        help='JSON string of node attributes. E.g., \'{"fontname": "Roboto Mono", "fontsize": "10"}\'',
+    )
+    p.add_argument(
+        "-e",
+        "--edge-attr",
+        type=lambda x: json.loads(x) if x else {},
+        default="{}",
+        help='JSON string of edge attributes. E.g., \'{"fontsize": "8"}\'',
+    )
 
     args = p.parse_args()
 
@@ -191,4 +206,10 @@ def draw(graph, path=None, format="png", verbose: bool = False):
             structs=ibis.struct({"a": [1, 2, 3], "b": {"c": 1, "d": 2}}),
         )
     )
-    expr.visualize(verbose=args.verbose > 0, label_edges=args.label_edges)
+
+    expr.visualize(
+        verbose=args.verbose > 0,
+        label_edges=args.label_edges,
+        node_attr=args.node_attr,
+        edge_attr=args.edge_attr,
+    )