Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for in-memory diagram generation using BytesIO #999

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add support for in-memory diagram generation using BytesIO #999

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 15 additions & 18 deletions diagrams/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
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.
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -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 = "",
Expand All @@ -89,6 +81,7 @@ def __init__(
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.

Expand All @@ -103,7 +96,7 @@ def __init__(
: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 = {}
Expand Down Expand Up @@ -152,18 +145,16 @@ def __init__(

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)
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 not self.output_buffer:
os.remove(self.filename)
setdiagram(None)

def _repr_png_(self):
Expand Down Expand Up @@ -193,9 +184,15 @@ def subgraph(self, dot: Digraph) -> None:
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:
Expand Down