Updating to allow context-manager invocations to work with new buildD…

…ot refactoring
hyakuhei · May 16, 2021 · c8c0691 · c8c0691
1 parent a9b94cc
commit c8c0691
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/attacktree/renderer.py b/attacktree/renderer.py
@@ -144,7 +144,11 @@ def render(
         fout: str = "png",
         renderOnExit=False,
     ):
+
         self.renderOnExit = renderOnExit
+        if root is None and self.root is not None:
+            root = self.root # sometimes we invoke as a context manager
+
         dot = self.buildDot(
             root=root, includeUnimplemented=renderUnimplemented, style=style
         )