From 5e07d5b11c979fdd4c737ea8778d9b914eb0915e Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 26 Sep 2026 22:14:42 +0800 Subject: [PATCH] fix(diagram): clear context when rendering fails --- diagrams/__init__.py | 10 ++++++---- tests/test_diagram.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 1fa343b0..f5b6f0bc 100644 --- a/diagrams/__init__.py +++ b/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") diff --git a/tests/test_diagram.py b/tests/test_diagram.py index 75fc02b8..01797543 100644 --- a/tests/test_diagram.py +++ b/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("<>") + 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")