From 9bb8e04b432890e3c09ccc215f3ab4c20ed0e136 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_image

Writes an image representation of a graph to a file. Several
common formats supported.
---
 .github/workflows/ci.yml         |  3 ++
 dev_requirements.txt             |  1 +
 src/graph_scheduler/utilities.py | 58 ++++++++++++++++++++++++++++++--
 tests/test_utilities.py          | 11 ++++++
 4 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d2bacf5d..e1f62831 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -58,6 +58,9 @@ jobs:
       if: ${{ matrix.python-version == 3.9 }}
       run: python -m pip install coveragepy-lcov "coverage<7.0.0"
 
+    - name: Install Graphviz
+      uses: tlylt/install-graphviz@v1
+
     - name: Lint with flake8
       shell: bash
       run: |
diff --git a/dev_requirements.txt b/dev_requirements.txt
index 39a462af..d801a57f 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -1,4 +1,5 @@
 psyneulink >=0.9.1.0
+pydot
 pytest
 pytest-benchmark
 pytest-cov
diff --git a/src/graph_scheduler/utilities.py b/src/graph_scheduler/utilities.py
index b80f3993..0deedb32 100644
--- a/src/graph_scheduler/utilities.py
+++ b/src/graph_scheduler/utilities.py
@@ -2,14 +2,14 @@
 import inspect
 import logging
 import weakref
-from typing import Dict, Hashable, Set
+from typing import Dict, Hashable, Set, Union
 
 import networkx as nx
 
 __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_image',
 ]
 
 
@@ -199,3 +199,57 @@ def networkx_digraph_to_dependency_dict(
                 res_graph[rec] = set()
             res_graph[rec].add(sender)
     return res_graph
+
+
+def output_graph_image(
+    graph: Union[typing_graph_dependency_dict, nx.Graph],
+    filename: str = None,
+    format: str = 'png',
+):
+    """
+    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. Defaults
+            to 'graph-scheduler-figure-<graph id>.<format>' in the current
+            directory.
+        format (str, optional): image format. Many common formats
+            supported. Pass None to display supported formats. Defaults
+            to png.
+
+    Requires:
+        - system graphviz: https://graphviz.org/download
+        - Python pydot: pip install pydot
+    """
+    if filename is None:
+        filename = f'graph-scheduler-figure-{id(graph)}.{format}'
+
+    if not isinstance(graph, nx.Graph):
+        graph = dependency_dict_to_networkx_digraph(graph)
+
+    try:
+        pd = nx.drawing.nx_pydot.to_pydot(graph)
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(
+            'Python pydot is required for output_graph_image.'
+            ' Install it with: pip install pydot'
+        ) from e
+
+    try:
+        pd.write(filename, format=format)
+    except AssertionError as e:
+        raise AssertionError(
+            f"Format '{format}' not recognized. Supported formats:"
+            f" {', '.join(pd.formats)}"
+        ) from e
+    except FileNotFoundError as e:
+        if '"dot" not found in path' in str(e):
+            raise FileNotFoundError(
+                'System graphviz is required for output_graph_image.'
+                ' Install it from https://graphviz.org/download'
+            ) from e
+        else:
+            raise
+
+    print(f'graph_scheduler.output_graph_image: wrote {format} to {filename}')
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index 4eb0e998..22bf8292 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -115,3 +115,14 @@ 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)
+@pytest.mark.parametrize('nx_type', nx_digraph_types)
+@pytest.mark.parametrize('format', ['png', 'jpg', 'svg'])
+def test_output_graph_image(graph, nx_type, format, tmp_path):
+    fname = f'{tmp_path}_fig.{format}'
+    nx_graph = graph_as_nx_graph(graph, nx_type)
+
+    gs.output_graph_image(graph, fname, format)
+    gs.output_graph_image(nx_graph, fname, format)