From 1ad084c48675b3f5f90d97fa5dfc720f52efe1e5 Mon Sep 17 00:00:00 2001
From: Cole Miles <colemiles132@gmail.com>
Date: Thu, 5 Sep 2024 20:51:23 +1200
Subject: [PATCH] Add save_to_disk flag to control diagram file saving

---
 diagrams/__init__.py | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index 6fe2a802..5719047d 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -89,6 +89,7 @@ class Diagram:
         graph_attr: Optional[dict] = None,
         node_attr: Optional[dict] = None,
         edge_attr: Optional[dict] = None,
+        save_to_disk: bool = True,  # Add flag to control saving behavior
     ):
         """Diagram represents a global diagrams context.
 
@@ -118,6 +119,9 @@ class Diagram:
             filename = "_".join(self.name.split()).lower()
         self.filename = filename
         self.dot = Digraph(self.name, filename=self.filename, strict=strict)
+        if save_to_disk is None:
+            save_to_disk = True
+
 
         # Set attributes.
         for k, v in self._default_graph_attrs.items():
@@ -152,6 +156,7 @@ class Diagram:
 
         self.show = show
         self.autolabel = autolabel
+        self.save_to_disk = save_to_disk
 
     def __str__(self) -> str:
         return str(self.dot)
@@ -161,11 +166,13 @@ 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)
+        if self.save_to_disk:  # Only render if save_to_disk is True
+            self.render()
+            # Remove the graphviz file leaving only the image.
+            os.remove(self.filename)
         setdiagram(None)
 
+
     def _repr_png_(self):
         return self.dot.pipe(format="png")
 
@@ -191,11 +198,12 @@ class Diagram:
         self.dot.subgraph(dot)
 
     def render(self) -> None:
-        if isinstance(self.outformat, list):
-            for one_format in self.outformat:
-                self.dot.render(format=one_format, view=self.show, quiet=True)
-        else:
-            self.dot.render(format=self.outformat, view=self.show, quiet=True)
+        if self.save_to_disk:  # Only save if save_to_disk is True
+            if isinstance(self.outformat, list):
+                for one_format in self.outformat:
+                    self.dot.render(format=one_format, view=self.show, quiet=True)
+            else:
+                self.dot.render(format=self.outformat, view=self.show, quiet=True)
 
 
 class Cluster: