pull/1262/merge
hhang 5 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 return self
def __exit__(self, exc_type, exc_value, traceback): def __exit__(self, exc_type, exc_value, traceback):
self.render() try:
# Remove the graphviz file leaving only the image. self.render()
os.remove(self.filename) # Remove the graphviz file leaving only the image.
setdiagram(None) os.remove(self.filename)
finally:
setdiagram(None)
def _repr_png_(self): def _repr_png_(self):
return self.dot.pipe(format="png") return self.dot.pipe(format="png")

@ -1,8 +1,11 @@
import os import os
import pathlib import pathlib
import shutil import shutil
import tempfile
import unittest import unittest
from graphviz import CalledProcessError
from diagrams import Cluster, Diagram, Edge, Node, getcluster, getdiagram, setcluster, setdiagram from diagrams import Cluster, Diagram, Edge, Node, getcluster, getdiagram, setcluster, setdiagram
@ -74,6 +77,15 @@ class DiagramTest(unittest.TestCase):
with self.assertRaises(EnvironmentError): with self.assertRaises(EnvironmentError):
Node("node") 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): def test_node_to_node(self):
with Diagram(name=os.path.join(self.name, "node_to_node"), show=False): with Diagram(name=os.path.join(self.name, "node_to_node"), show=False):
node1 = Node("node1") node1 = Node("node1")

Loading…
Cancel
Save