pull/1262/merge
hhang 4 days ago committed by GitHub
commit e566ca39e7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -228,10 +228,12 @@ class Diagram:
return self
def __exit__(self, exc_type, exc_value, traceback):
self.render()
# Remove the graphviz file leaving only the image.
os.remove(self.filename)
setdiagram(None)
try:
self.render()
# Remove the graphviz file leaving only the image.
os.remove(self.filename)
finally:
setdiagram(None)
def _repr_png_(self):
return self.dot.pipe(format="png")

@ -1,8 +1,11 @@
import os
import pathlib
import shutil
import tempfile
import unittest
from graphviz import CalledProcessError
from diagrams import Cluster, Diagram, Edge, Node, getcluster, getdiagram, setcluster, setdiagram
@ -74,6 +77,15 @@ class DiagramTest(unittest.TestCase):
with self.assertRaises(EnvironmentError):
Node("node")
def test_render_error_clears_global_context(self):
with tempfile.TemporaryDirectory() as directory:
with self.assertRaises(CalledProcessError):
with Diagram(filename=os.path.join(directory, "broken"), show=False):
Node("<<invalid>>")
self.assertIsNone(getdiagram())
with self.assertRaises(EnvironmentError):
Node("outside diagram")
def test_node_to_node(self):
with Diagram(name=os.path.join(self.name, "node_to_node"), show=False):
node1 = Node("node1")

Loading…
Cancel
Save