Fix cyclical import of plots (mckinsey#152)

Co-authored-by: philip_pilgerstorfer <philip.pilgerstorfer!@quantumblack.com>
vishalbelsare · Apr 12, 2021 · 88f23cd · 88f23cd
1 parent 0ed3bdc
commit 88f23cd
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -31,7 +31,7 @@ repos:
  exclude: ^causalnex/ebaybbn
 
 - repo: https://github.com/psf/black
- rev: stable
+ rev: 20.8b1
  hooks:
  - id: black
 

diff --git a/RELEASE.md b/RELEASE.md
@@ -1,5 +1,7 @@
 # Upcoming release
 
+* Fixes cyclical import of `causalnex.plots`, as per #106.
+
 # Release 0.9.2
 * Remove Boston housing dataset from "sklearn tutorial", see #91 for more information.
 * Update pylint version to 2.7

diff --git a/causalnex/plots/plots.py b/causalnex/plots/plots.py
@@ -33,11 +33,9 @@
 
 import networkx as nx
 
-from causalnex.structure.structuremodel import StructureModel
-
 
 def plot_structure(
- sm: StructureModel,
+ sm: nx.DiGraph,
  prog: str = "neato",
  all_node_attributes: Dict[str, str] = None,
  all_edge_attributes: Dict[str, str] = None,
@@ -187,12 +185,12 @@ def color_gradient_string(from_color: str, to_color: str, steps: int) -> str:
 
 
 def _add_attributes(
- sm: StructureModel,
+ sm: nx.DiGraph,
  all_node_attributes: Dict[str, str] = None,
  all_edge_attributes: Dict[str, str] = None,
  node_attributes: Dict[str, Dict[str, str]] = None,
  edge_attributes: Dict[str, Dict[str, str]] = None,
-) -> StructureModel:
+) -> nx.DiGraph:
  _sm = deepcopy(sm)
 
  # shift labels to be above nodes