Browse Source

Add save_to_disk flag to control diagram file saving

pull/1017/head
Cole Miles 2 months ago
parent
commit
1ad084c486
1 changed files with 16 additions and 8 deletions
  1. +16
    -8
      diagrams/__init__.py

+ 16
- 8
diagrams/__init__.py View File

@@ -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:


Loading…
Cancel
Save