Skip to content

Commit

Permalink
feat(output): Generate dot file and support multi outformat.(mingramm…
Browse files Browse the repository at this point in the history
…er#441) Re (mingrammer#592)

* feat(output): Generate dot file and support multi outformat.(mingrammer#441)

* [fix] forget to clean the dot generated file.

* [fix] indentation

* [fix] Review + add more cases in unittest

* [fix] Add dot in the test
  • Loading branch information
gabriel-tessier authored Feb 9, 2022
1 parent 1422de9 commit 5055e00
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 26 deletions.
40 changes: 18 additions & 22 deletions diagrams/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def setcluster(cluster):
class Diagram:
__directions = ("TB", "BT", "LR", "RL")
__curvestyles = ("ortho", "curved")
__outformats = ("png", "jpg", "svg", "pdf")
__outformats = ("png", "jpg", "svg", "pdf", "dot")

# fmt: off
_default_graph_attrs = {
Expand Down Expand Up @@ -127,8 +127,13 @@ def __init__(
raise ValueError(f'"{curvestyle}" is not a valid curvestyle')
self.dot.graph_attr["splines"] = curvestyle

if not self._validate_outformat(outformat):
raise ValueError(f'"{outformat}" is not a valid output format')
if isinstance(outformat, list):
for one_format in outformat:
if not self._validate_outformat(one_format):
raise ValueError(f'"{one_format}" is not a valid output format')
else:
if not self._validate_outformat(outformat):
raise ValueError(f'"{outformat}" is not a valid output format')
self.outformat = outformat

# Merge passed in attributes
Expand All @@ -155,25 +160,13 @@ def _repr_png_(self):
return self.dot.pipe(format="png")

def _validate_direction(self, direction: str) -> bool:
direction = direction.upper()
for v in self.__directions:
if v == direction:
return True
return False
return direction.upper() in self.__directions

def _validate_curvestyle(self, curvestyle: str) -> bool:
curvestyle = curvestyle.lower()
for v in self.__curvestyles:
if v == curvestyle:
return True
return False
return curvestyle.lower() in self.__curvestyles

def _validate_outformat(self, outformat: str) -> bool:
outformat = outformat.lower()
for v in self.__outformats:
if v == outformat:
return True
return False
return outformat.lower() in self.__outformats

def node(self, nodeid: str, label: str, **attrs) -> None:
"""Create a new node."""
Expand All @@ -188,7 +181,11 @@ def subgraph(self, dot: Digraph) -> None:
self.dot.subgraph(dot)

def render(self) -> None:
self.dot.render(format=self.outformat, view=self.show, quiet=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:
Expand Down Expand Up @@ -263,9 +260,8 @@ def __exit__(self, exc_type, exc_value, traceback):

def _validate_direction(self, direction: str):
direction = direction.upper()
for v in self.__directions:
if v == direction:
return True
if direction in self.__directions:
return True
return False

def node(self, nodeid: str, label: str, **attrs) -> None:
Expand Down
12 changes: 11 additions & 1 deletion docs/guides/diagram.md
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ diag

You can specify the output file format with `outformat` parameter. Default is **png**.

> (png, jpg, svg, and pdf) are allowed.
> (png, jpg, svg, pdf and dot) are allowed.
```python
from diagrams import Diagram
Expand All @@ -54,6 +54,16 @@ with Diagram("Simple Diagram", outformat="jpg"):
EC2("web")
```

The `outformat` parameter also support list to output all the defined output in one call.

```python
from diagrams import Diagram
from diagrams.aws.compute import EC2

with Diagram("Simple Diagram Multi Output", outformat=["jpg", "png", "dot"]):
EC2("web")
```

You can specify the output filename with `filename` parameter. The extension shouldn't be included, it's determined by the `outformat` parameter.

```python
Expand Down
18 changes: 15 additions & 3 deletions tests/test_diagram.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def tearDown(self):

def test_validate_direction(self):
# Normal directions.
for dir in ("TB", "BT", "LR", "RL"):
for dir in ("TB", "BT", "LR", "RL", "tb"):
Diagram(direction=dir)

# Invalid directions.
Expand All @@ -36,7 +36,7 @@ def test_validate_direction(self):

def test_validate_curvestyle(self):
# Normal directions.
for cvs in ("ortho", "curved"):
for cvs in ("ortho", "curved", "CURVED"):
Diagram(curvestyle=cvs)

# Invalid directions.
Expand All @@ -46,7 +46,7 @@ def test_validate_curvestyle(self):

def test_validate_outformat(self):
# Normal output formats.
for fmt in ("png", "jpg", "svg", "pdf"):
for fmt in ("png", "jpg", "svg", "pdf", "PNG", "dot"):
Diagram(outformat=fmt)

# Invalid output formats.
Expand Down Expand Up @@ -108,6 +108,18 @@ def test_empty_name(self):
Node("node1")
self.assertTrue(os.path.exists(f"{self.name}.png"))

def test_outformat_list(self):
"""Check that outformat render all the files from the list."""
self.name = 'diagrams_image'
with Diagram(show=False, outformat=["dot", "png"]):
Node("node1")
# both files must exist
self.assertTrue(os.path.exists(f"{self.name}.png"))
self.assertTrue(os.path.exists(f"{self.name}.dot"))

# clean the dot file as it only generated here
os.remove(self.name + ".dot")


class ClusterTest(unittest.TestCase):
def setUp(self):
Expand Down

0 comments on commit 5055e00

Please sign in to comment.