From f06d8f3b8b8127abe494aef2333dd12e2486d35c Mon Sep 17 00:00:00 2001
From: Andrechang <andrechang67@gmail.com>
Date: Fri, 15 Oct 2021 22:19:44 -0700
Subject: [PATCH] add build graph networkx 15 Oct 2021

build hiddelayer as networkx graph
---
 hiddenlayer/graph.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/hiddenlayer/graph.py b/hiddenlayer/graph.py
index d8adbd3..1f15f75 100644
--- a/hiddenlayer/graph.py
+++ b/hiddenlayer/graph.py
@@ -354,6 +354,33 @@ def build_dot(self):
             dot.edge(str(a), str(b), label)
         return dot
 
+    def build_nx(self):
+        """Generate a Networkx graph.
+
+        Returns a Networkx Digraph object.
+        """
+        import networkx as nx
+
+        # Build GraphViz Digraph
+        G = nx.DiGraph()
+        nodes = []
+        for k, n in self.nodes.items():
+            label = "{} ".format(n.title)
+            if n.caption:
+                label += "{} ".format(n.caption)
+            if n.repeat > 1:
+                label += "x{}".format(n.repeat)
+            label = " " + label + " "
+            nodes.append((str(k), {'label': label, 'shape': 'box'}))
+        G.add_nodes_from(nodes)
+
+        for a, b, label in self.edges:
+            if isinstance(label, (list, tuple)):
+                label = "x".join([str(l or "?") for l in label])
+            G.add_edge(str(a), str(b), label=label)
+            # dot.edge(str(a), str(b), label)
+        return G
+
     def _repr_svg_(self):
         """Allows Jupyter notebook to render the graph automatically."""
         return self.build_dot()._repr_svg_()