diff --git a/.gitignore b/.gitignore index 784dc50..80775bb 100644 --- a/.gitignore +++ b/.gitignore @@ -90,4 +90,4 @@ local_settings.py .env db.sqlite3 /local/ -.vscode +.vscode/ diff --git a/mermaid/flowchart/__init__.py b/mermaid/flowchart/__init__.py index d59d63e..7e676ea 100644 --- a/mermaid/flowchart/__init__.py +++ b/mermaid/flowchart/__init__.py @@ -8,15 +8,17 @@ class FlowChart(Graph): def __init__(self, title: str, nodes: list[Node] = None, - links: list[Link] = None) -> None: + links: list[Link] = None, + orientation: str = 'TB') -> None: super().__init__(title, '') + self.orientation: str = orientation self.nodes: list[Node] = nodes if nodes is not None else [] self.links: list[Link] = links if links is not None else [] self._build_script() def _build_script(self) -> None: super()._build_script() - script: str = '\nflowchart' + script: str = f'\nflowchart {self.orientation}' for node in self.nodes: script += f'\n\t{node}' for link in self.links: diff --git a/mermaid/tests/test_flowchart.py b/mermaid/tests/test_flowchart.py index dc815fa..f43d97b 100644 --- a/mermaid/tests/test_flowchart.py +++ b/mermaid/tests/test_flowchart.py @@ -85,7 +85,7 @@ def test_str_link_with_no_default_value(self): class TestFlowChart(unittest.TestCase): - def test_make_flowchart_script(self): + def test_make_flowchart_script_with_default_orientation(self): nodes = [Node('First Node'), Node('Second Node'), Node('Third Node')] links = [ @@ -96,7 +96,30 @@ def test_make_flowchart_script(self): expect_script: str = f"""--- title: simple flowchart --- -flowchart +flowchart TB +\t{nodes[0]} +\t{nodes[1]} +\t{nodes[2]} +\t{links[0]} +\t{links[1]} +""" + self.assertEqual(expect_script, flowchart.script) + + def test_make_flowchart_script_without_default_orientation(self): + + nodes = [Node('First Node'), Node('Second Node'), Node('Third Node')] + links = [ + Link(nodes[0], nodes[1], head_left='cross'), + Link(nodes[1], nodes[2], head_right='bullet') + ] + flowchart: FlowChart = FlowChart('simple flowchart', + nodes, + links, + orientation='LR') + expect_script: str = f"""--- +title: simple flowchart +--- +flowchart LR \t{nodes[0]} \t{nodes[1]} \t{nodes[2]}