fix(diagram): clear context when rendering fails

diagrams/__init__.py

@@ -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")

tests/test_diagram.py

@@ -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")