Skip to content

Commit

Permalink
feat(option): add curvestyle option (ortho or curved)
Browse files Browse the repository at this point in the history
  • Loading branch information
mingrammer committed Jun 28, 2020
1 parent 91a2910 commit f8ac547
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 0 deletions.
14 changes: 14 additions & 0 deletions diagrams/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ def setcluster(cluster):

class Diagram:
__directions = ("TB", "BT", "LR", "RL")
__curvestyles = ("ortho", "curved")
__outformats = ("png", "jpg", "svg", "pdf")

# fmt: off
Expand Down Expand Up @@ -78,6 +79,7 @@ def __init__(
name: str = "",
filename: str = "",
direction: str = "LR",
curvestyle: str = "ortho",
outformat: str = "png",
show: bool = True,
graph_attr: dict = {},
Expand All @@ -91,6 +93,7 @@ def __init__(
:param filename: The output filename, without the extension (.png).
If not given, it will be generated from the name.
:param direction: Data flow direction. Default is 'left to right'.
:param curvestyle: Curve bending style. One of "ortho" or "curved".
:param outformat: Output file format. Default is 'png'.
:param show: Open generated image after save if true, just only save otherwise.
:param graph_attr: Provide graph_attr dot config attributes.
Expand All @@ -117,6 +120,10 @@ def __init__(
raise ValueError(f'"{direction}" is not a valid direction')
self.dot.graph_attr["rankdir"] = direction

if not self._validate_curvestyle(curvestyle):
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')
self.outformat = outformat
Expand Down Expand Up @@ -151,6 +158,13 @@ def _validate_direction(self, direction: str) -> bool:
return True
return False

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

def _validate_outformat(self, outformat: str) -> bool:
outformat = outformat.lower()
for v in self.__outformats:
Expand Down
10 changes: 10 additions & 0 deletions tests/test_diagram.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,16 @@ def test_validate_direction(self):
with self.assertRaises(ValueError):
Diagram(direction=dir)

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

# Invalid directions.
for cvs in ("tangent", "unknown"):
with self.assertRaises(ValueError):
Diagram(curvestyle=cvs)

def test_validate_outformat(self):
# Normal output formats.
for fmt in ("png", "jpg", "svg", "pdf"):
Expand Down

0 comments on commit f8ac547

Please sign in to comment.