From 1c2c20a04ccd0ac434cc12b673e296bec1d53e38 Mon Sep 17 00:00:00 2001
From: Katherine Mantel <kmantel@princeton.edu>
Date: Mon, 6 Nov 2023 23:49:55 +0000
Subject: [PATCH] utilities: add output_graph_png

writes an image representation of a graph to a png file
---
 src/graph_scheduler/utilities.py | 25 ++++++++++++++++++++++++-
 tests/test_utilities.py          |  5 +++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/graph_scheduler/utilities.py b/src/graph_scheduler/utilities.py
index b80f399..c78ee0d 100644
--- a/src/graph_scheduler/utilities.py
+++ b/src/graph_scheduler/utilities.py
@@ -9,7 +9,7 @@
 __all__ = [
     'clone_graph', 'dependency_dict_to_networkx_digraph',
     'disable_debug_logging', 'enable_debug_logging',
-    'networkx_digraph_to_dependency_dict',
+    'networkx_digraph_to_dependency_dict', 'output_graph_png',
 ]
 
 
@@ -199,3 +199,26 @@ def networkx_digraph_to_dependency_dict(
                 res_graph[rec] = set()
             res_graph[rec].add(sender)
     return res_graph
+
+
+def output_graph_png(graph: typing_graph_dependency_dict, filename: str = None):
+    """
+    Writes an image representation of **graph** to file **filename**
+
+    Args:
+        graph: a graph in dependency dict form
+        filename (str, optional): full path of image to write. Appends
+            '.png' if not present. Defaults to
+            'graph-scheduler-figure-<graph id>.png' in the current
+            directory.
+    """
+    ext = '.png'
+    if filename is None:
+        filename = f'graph-scheduler-figure-{id(graph)}{ext}'
+    elif not filename.endswith(ext):
+        filename += ext
+
+    nx_graph = dependency_dict_to_networkx_digraph(graph)
+    pd = nx.drawing.nx_pydot.to_pydot(nx_graph)
+    pd.write_png(filename)
+    print(f'graph_scheduler.output_graph_png: wrote {ext[1:]} to {filename}')
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 4eb0e99..558ca21 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -115,3 +115,8 @@ def test_convert_from_nx_graph(graph, nx_type):
     )
     assert nx_graph.nodes == nx_graph.nodes
     assert nx_graph.edges == res.edges
+
+
+@pytest.mark.parametrize('graph', test_graphs)
+def test_output_graph_png(graph, tmp_path):
+    gs.output_graph_png(graph, f'{tmp_path}_fig.png')