Drop hard networkx dependency

stitchfix · Dec 12, 2021 · f00dbf4 · f00dbf4
1 parent d17a5ca
commit f00dbf4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/hamilton/graph.py b/hamilton/graph.py
@@ -8,8 +8,6 @@
 from types import ModuleType
 from typing import Type, Dict, Any, Callable, Tuple, Set, Collection, List
 
-import networkx
-
 from hamilton import function_modifiers
 from hamilton import node
 from hamilton.node import NodeSource, DependencyType
@@ -147,6 +145,14 @@ def display(nodes: Set[node.Node], user_nodes: Set[node.Node], output_file_path:
                 '\n\n  pip install graphviz'
             )
 
+        try:
+            import networkx
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                'networkx is required for visualizing the function graph. Install it with:'
+                '\n\n  pip install networkx'
+            )
+
         dot = graphviz.Digraph(comment='Dependency Graph')
 
         for n in nodes:

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,3 @@
 dataclasses
-networkx
 numpy
 pandas