From 29a56bc3595ebc031a02686d49458b5d8d8c928c Mon Sep 17 00:00:00 2001
From: rithvikgrandhi <grandhirithwik@gmail.com>
Date: Mon, 20 May 2024 15:57:20 +0530
Subject: [PATCH] Add support for in-memory diagram generation using BytesIO

---
 diagrams/__init__.py | 33 +++++++++++++++------------------
 1 file changed, 15 insertions(+), 18 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index 6fe2a80..b1e141f 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -3,7 +3,7 @@ import os
 import uuid
 from pathlib import Path
 from typing import Dict, List, Optional, Union
-
+from io import BytesIO
 from graphviz import Digraph
 
 # Global contexts for a diagrams and a cluster.
@@ -42,7 +42,7 @@ class Diagram:
     __curvestyles = ("ortho", "curved")
     __outformats = ("png", "jpg", "svg", "pdf", "dot")
 
-    # fmt: off
+    # Default graph attributes
     _default_graph_attrs = {
         "pad": "2.0",
         "splines": "ortho",
@@ -59,10 +59,6 @@ class Diagram:
         "width": "1.4",
         "height": "1.4",
         "labelloc": "b",
-        # imagepos attribute is not backward compatible
-        # TODO: check graphviz version to see if "imagepos" is available >= 2.40
-        # https://github.com/xflr6/graphviz/blob/master/graphviz/backend.py#L248
-        # "imagepos": "tc",
         "imagescale": "true",
         "fontname": "Sans-Serif",
         "fontsize": "13",
@@ -72,10 +68,6 @@ class Diagram:
         "color": "#7B8894",
     }
 
-    # fmt: on
-
-    # TODO: Label position option
-    # TODO: Save directory option (filename + directory?)
     def __init__(
         self,
         name: str = "",
@@ -89,6 +81,7 @@ class Diagram:
         graph_attr: Optional[dict] = None,
         node_attr: Optional[dict] = None,
         edge_attr: Optional[dict] = None,
+        output_buffer: Optional[BytesIO] = None  # New parameter
     ):
         """Diagram represents a global diagrams context.
 
@@ -103,7 +96,7 @@ class Diagram:
         :param graph_attr: Provide graph_attr dot config attributes.
         :param node_attr: Provide node_attr dot config attributes.
         :param edge_attr: Provide edge_attr dot config attributes.
-        :param strict: Rendering should merge multi-edges.
+        :param output_buffer: Optional in-memory buffer to write the output to.
         """
         if graph_attr is None:
             graph_attr = {}
@@ -152,9 +145,7 @@ class Diagram:
 
         self.show = show
         self.autolabel = autolabel
-
-    def __str__(self) -> str:
-        return str(self.dot)
+        self.output_buffer = output_buffer  # Store the buffer
 
     def __enter__(self):
         setdiagram(self)
@@ -162,8 +153,8 @@ class Diagram:
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.render()
-        # Remove the graphviz file leaving only the image.
-        os.remove(self.filename)
+        if not self.output_buffer:
+            os.remove(self.filename)
         setdiagram(None)
 
     def _repr_png_(self):
@@ -193,9 +184,15 @@ class Diagram:
     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)
+                if self.output_buffer:
+                    self.output_buffer.write(self.dot.pipe(format=one_format))
+                else:
+                    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.output_buffer:
+                self.output_buffer.write(self.dot.pipe(format=self.outformat))
+            else:
+                self.dot.render(format=self.outformat, view=self.show, quiet=True)
 
 
 class Cluster: