Add option to render a subgraph of a hugr.

Author: cqc-alec

hugr-py/src/hugr/hugr/base.py

@@ -1152,21 +1152,32 @@ def load_json(cls, json_str: str) -> Hugr:
         serial = SerialHugr.load_json(json_dict)
         return cls._from_serial(serial)
 
-    def render_dot(self, config: RenderConfig | None = None) -> gv.Digraph:
+    def render_dot(
+        self, config: RenderConfig | None = None, root: Node | None = None
+    ) -> gv.Digraph:
         """Render the HUGR to a graphviz Digraph.
 
         Args:
             config: Render configuration.
+            root: Root node defining the set of nodes to render. By default this is the
+                module root and all nodes are rendered. If this is a container node, all
+                nodes under it are rendered. Every incoming edge to the rendered set and
+                outgoing edge from it is also shown, with its other endpoint labelled
+                with its node index.
 
         Returns:
             The graphviz Digraph.
         """
         from .render import DotRenderer
 
-        return DotRenderer(config).render(self)
+        return DotRenderer(config).render(self, root)
 
     def store_dot(
-        self, filename: str, format: str = "svg", config: RenderConfig | None = None
+        self,
+        filename: str,
+        format: str = "svg",
+        config: RenderConfig | None = None,
+        root: Node | None = None,
     ) -> None:
         """Render the HUGR to a graphviz dot file.
 
@@ -1175,7 +1186,12 @@ def store_dot(
             format: The format used for rendering ('pdf', 'png', etc.).
                 Defaults to SVG.
             config: Render configuration.
+            root: Root node defining the set of nodes to render. By default this is the
+                module root and all nodes are rendered. If this is a container node, all
+                nodes under it are rendered. Every incoming edge to the rendered set and
+                outgoing edge from it is also shown, with its other endpoint labelled
+                with its node index.
         """
         from .render import DotRenderer
 
-        DotRenderer(config).store(self, filename=filename, format=format)
+        DotRenderer(config).store(self, filename=filename, format=format, root=root)

hugr-py/src/hugr/hugr/render.py

@@ -92,42 +92,60 @@ class DotRenderer:
 
     def __init__(self, config: RenderConfig | None = None) -> None:
         self.config = config or RenderConfig()
+        self.nodes = set()
 
-    def render(self, hugr: Hugr) -> Digraph:
-        """Render a HUGR to a graphviz dot object."""
+    def render(self, hugr: Hugr, root: Node | None = None) -> Digraph:
+        """Render a HUGR to a graphviz dot object.
+
+        Args:
+            hugr: The HUGR to render.
+            root: Root node defining the set of nodes to render. By default this is the
+                module root and all nodes are rendered. If this is a container node, all
+                nodes under it are rendered. Every incoming edge to the rendered set and
+                outgoing edge from it is also shown, with its other endpoint labelled
+                with its node index.
+        """
+        root = root or hugr.module_root
         graph_attr = {
             "rankdir": "",
             "ranksep": "0.1",
             "nodesep": "0.15",
             "margin": "0",
             "bgcolor": self.config.palette.background,
         }
-        if name := hugr[hugr.module_root].metadata.get("name", None):
+        if name := hugr[root].metadata.get("name", None):
             name = html.escape(str(name))
         else:
             name = ""
 
         graph = gv.Digraph(name, strict=False)
         graph.attr(**graph_attr)
 
-        self._viz_node(hugr.module_root, hugr, graph)
+        self._viz_node(root, hugr, graph)
 
         for src_port, tgt_port in hugr.links():
             kind = hugr.port_kind(src_port)
             self._viz_link(src_port, tgt_port, kind, graph)
 
         return graph
 
-    def store(self, hugr: Hugr, filename: str, format: str = "svg") -> None:
+    def store(
+        self, hugr: Hugr, filename: str, format: str = "svg", root: Node | None = None
+    ) -> None:
         """Render a HUGR and save it to a file.
 
         Args:
             hugr: The HUGR to render.
             filename: Filename for saving the rendered graph.
             format: The format used for rendering ('pdf', 'png', etc.).
                 Defaults to SVG.
+            root: Root node defining the set of nodes to render. By default this is the
+                module root and all nodes are rendered. If this is a container node, all
+                nodes under it are rendered. Every incoming edge to the rendered set and
+                outgoing edge from it is also shown, with its other endpoint labelled
+                with its node index.
         """
-        gv_graph = self.render(hugr)
+        gv_graph = self.render(hugr, root=root)
         gv_graph.render(filename, format=format)
 
     _FONTFACE = "monospace"
@@ -266,6 +284,7 @@ def _viz_node(self, node: Node, hugr: Hugr, graph: Digraph) -> None:
                     self._viz_node(child, hugr, sub)
                 html_label = self._format_html_label(**label_config)
                 sub.node(f"{node.idx}", shape="plain", label=f"<{html_label}>")
+                self.nodes.add(node)
                 sub.attr(
                     label="",
                     margin="10",
@@ -275,6 +294,7 @@ def _viz_node(self, node: Node, hugr: Hugr, graph: Digraph) -> None:
         else:
             html_label = self._format_html_label(**label_config)
             graph.node(f"{node.idx}", label=f"<{html_label}>", shape="plain")
+            self.nodes.add(node)
 
     def _viz_link(
         self, src_port: OutPort, tgt_port: InPort, kind: Kind, graph: Digraph
@@ -302,10 +322,11 @@ def _viz_link(
             case _:
                 assert_never(kind)
 
-        graph.edge(
-            self._out_port_name(src_port),
-            self._in_port_name(tgt_port),
-            label=label,
-            color=color,
-            **edge_attr,
-        )
+        if src_port.node in self.nodes or tgt_port.node in self.nodes:
+            graph.edge(
+                self._out_port_name(src_port),
+                self._in_port_name(tgt_port),
+                label=label,
+                color=color,
+                **edge_attr,
+            )