Merge pull request #94 from pyiron/draw_more

Expose a subset of graphviz.Graph.render options inside Node.draw

pyiron_workflow/node.py

@@ -25,6 +25,8 @@
 from pyiron_workflow.util import SeabornColors
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     import graphviz
 
     from pyiron_workflow.channels import Channel
@@ -587,10 +589,26 @@ def color(self) -> str:
         return SeabornColors.white
 
     def draw(
-        self, depth: int = 1, rankdir: Literal["LR", "TB"] = "LR"
+        self,
+        depth: int = 1,
+        rankdir: Literal["LR", "TB"] = "LR",
+        save: bool = False,
+        view: bool = False,
+        directory: Optional[Path | str] = None,
+        filename: Optional[Path | str] = None,
+        format: Optional[str] = None,
+        cleanup: bool = True,
     ) -> graphviz.graphs.Digraph:
         """
-        Draw the node structure.
+        Draw the node structure and return it as a graphviz object.
+
+        A selection of the `graphviz.Graph.render` method options are exposed, and if
+        `view` or `filename` is provided, this will be called before returning the
+        graph.
+        The graph file and rendered image will be stored in the node's working
+        directory.
+        This is purely for convenience -- since we directly return a graphviz object
+        you can instead use this to leverage the full power of graphviz.
 
         Args:
             depth (int): How deeply to decompose the representation of composite nodes
@@ -600,17 +618,37 @@ def draw(
                 max depth of the node will have no adverse side effects.
             rankdir ("LR" | "TB"): Use left-right or top-bottom graphviz `rankdir` to
                 orient the flow of the graph.
+            save (bool): Render the graph image. (Default is False. When True, all
+                other defaults will yield a PDF in the node's working directory.)
+            view (bool): `graphviz.Graph.render` argument, open the rendered result
+                with the default application. (Default is False. When True, default
+                values for the directory and filename are supplied by the node working
+                directory and label.)
+            directory (Path|str|None): `graphviz.Graph.render` argument, (sub)directory
+                for source saving and rendering. (Default is None, which uses the
+                node's working directory.)
+            filename (Path|str): `graphviz.Graph.render` argument, filename for saving
+                the source. (Default is None, which uses the node label + `"_graph"`.
+            format (str|None): `graphviz.Graph.render` argument, the output format used
+                for rendering ('pdf', 'png', etc.).
+            cleanup (bool): `graphviz.Graph.render` argument, delete the source file
+                after successful rendering. (Default is True -- unlike graphviz.)
 
         Returns:
             (graphviz.graphs.Digraph): The resulting graph object.
-
-        Note:
-            The graphviz docs will elucidate all the possibilities of what to do with
-            the returned object, but the thing you are most likely to need is the
-            `render` method, which allows you to save the resulting graph as an image.
-            E.g. `self.draw().render(filename="my_node", format="png")`.
         """
-        return GraphvizNode(self, depth=depth, rankdir=rankdir).graph
+        graph = GraphvizNode(self, depth=depth, rankdir=rankdir).graph
+        if save or view or filename is not None:
+            directory = self.working_directory.path if directory is None else directory
+            filename = self.label + "_graph" if filename is None else filename
+            graph.render(
+                view=view,
+                directory=directory,
+                filename=filename,
+                format=format,
+                cleanup=cleanup,
+            )
+        return graph
 
     def activate_strict_hints(self):
         """Enable type hint checks for all data IO"""

tests/unit/test_node.py

@@ -305,3 +305,33 @@ def test_working_directory(self):
             msg="Just want to make sure we cleaned up after ourselves"
         )
 
+    def test_draw(self):
+        try:
+            self.n1.draw()
+            self.assertFalse(
+                any(self.n1.working_directory.path.iterdir())
+            )
+
+            fmt = "pdf"  # This is just so we concretely know the filename suffix
+            self.n1.draw(save=True, format=fmt)
+            expected_name = self.n1.label + "_graph." + fmt
+            # That name is just an implementation detail, update it as needed
+            self.assertTrue(
+                self.n1.working_directory.path.joinpath(expected_name).is_file(),
+                msg="If `save` is called, expect the rendered image to exist in the working"
+                    "directory"
+            )
+
+            user_specified_name = "foo"
+            self.n1.draw(filename=user_specified_name, format=fmt)
+            expected_name = user_specified_name + "." + fmt
+            self.assertTrue(
+                self.n1.working_directory.path.joinpath(expected_name).is_file(),
+                msg="If the user specifies a filename, we should assume they want the "
+                    "thing saved"
+            )
+        finally:
+            # No matter what happens in the tests, clean up after yourself
+            self.n1.working_directory.delete()
+
+