Skip to content

Commit e426902

Browse files
kayjankayjan
authored
Merge pull request #131 from kayjan/fix-mermaid
Change: Mermaid delimiter, abstract assertions and constants to anoth…
2 parents 9df4f3f + 56ecb1a commit e426902

File tree

5 files changed

+178
-131
lines changed

5 files changed

+178
-131
lines changed

‎README.md‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -915,10 +915,10 @@ print(mermaid_md)
915915
```mermaid
916916
%%{ init: { 'flowchart': { 'curve': 'basis' } } }%%
917917
flowchart TB
918-
0("a") --> 00("b")
919-
00 --> 000("d")
920-
00 --> 001("e")
921-
0("a") --> 01("c")
918+
0("a") --> 0-0("b")
919+
0-0 --> 0-0-0("d")
920+
0-0 --> 0-0-1("e")
921+
0("a") --> 0-1("c")
922922
classDef default stroke-width:1
923923
```
924924

‎bigtree/tree/export.py‎

Lines changed: 12 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
from bigtree.node.node import Node
88
from bigtree.tree.search import find_path
9+
from bigtree.utils.assertions import assert_key_in_dict, assert_str_in_list
10+
from bigtree.utils.constants import ExportConstants, MermaidConstants
911
from bigtree.utils.exceptions import (
1012
optional_dependencies_image,
1113
optional_dependencies_pandas,
@@ -41,16 +43,6 @@
4143

4244
T = TypeVar("T", bound=Node)
4345

44-
available_styles = {
45-
"ansi": ("| ", "|-- ", "`-- "),
46-
"ascii": ("| ", "|-- ", "+-- "),
47-
"const": ("\u2502 ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 "),
48-
"const_bold": ("\u2503 ", "\u2523\u2501\u2501 ", "\u2517\u2501\u2501 "),
49-
"rounded": ("\u2502 ", "\u251c\u2500\u2500 ", "\u2570\u2500\u2500 "),
50-
"double": ("\u2551 ", "\u2560\u2550\u2550 ", "\u255a\u2550\u2550 "),
51-
"custom": ("", "", ""),
52-
}
53-
5446

5547
def print_tree(
5648
tree: T,
@@ -324,6 +316,7 @@ def yield_tree(
324316
style (str): style of print, defaults to abstract style
325317
custom_style (Iterable[str]): style of stem, branch and final stem, used when `style` is set to 'custom'
326318
"""
319+
available_styles = ExportConstants.AVAILABLE_STYLES
327320
if style not in available_styles.keys():
328321
raise ValueError(
329322
f"Choose one of {available_styles.keys()} style, use `custom` to define own style"
@@ -1001,66 +994,16 @@ def tree_to_mermaid(
1001994
"""
1002995
from bigtree.tree.helper import clone_tree
1003996

1004-
rankdirs = ["TB", "BT", "LR", "RL"]
1005-
line_shapes = [
1006-
"basis",
1007-
"bumpX",
1008-
"bumpY",
1009-
"cardinal",
1010-
"catmullRom",
1011-
"linear",
1012-
"monotoneX",
1013-
"monotoneY",
1014-
"natural",
1015-
"step",
1016-
"stepAfter",
1017-
"stepBefore",
1018-
]
1019-
node_shapes = {
1020-
"rounded_edge": """("{label}")""",
1021-
"stadium": """(["{label}"])""",
1022-
"subroutine": """[["{label}"]]""",
1023-
"cylindrical": """[("{label}")]""",
1024-
"circle": """(("{label}"))""",
1025-
"asymmetric": """>"{label}"]""",
1026-
"rhombus": """{{"{label}"}}""",
1027-
"hexagon": """{{{{"{label}"}}}}""",
1028-
"parallelogram": """[/"{label}"/]""",
1029-
"parallelogram_alt": """[\\"{label}"\\]""",
1030-
"trapezoid": """[/"{label}"\\]""",
1031-
"trapezoid_alt": """[\\"{label}"/]""",
1032-
"double_circle": """((("{label}")))""",
1033-
}
1034-
edge_arrows = {
1035-
"normal": "-->",
1036-
"bold": "==>",
1037-
"dotted": "-.->",
1038-
"open": "---",
1039-
"bold_open": "===",
1040-
"dotted_open": "-.-",
1041-
"invisible": "~~~",
1042-
"circle": "--o",
1043-
"cross": "--x",
1044-
"double_normal": "<-->",
1045-
"double_circle": "o--o",
1046-
"double_cross": "x--x",
1047-
}
997+
rankdirs = MermaidConstants.RANK_DIR
998+
line_shapes = MermaidConstants.LINE_SHAPES
999+
node_shapes = MermaidConstants.NODE_SHAPES
1000+
edge_arrows = MermaidConstants.EDGE_ARROWS
10481001

10491002
# Assertions
1050-
if rankdir not in rankdirs:
1051-
raise ValueError(f"Invalid input, check `rankdir` should be one of {rankdirs}")
1052-
if node_shape not in node_shapes:
1053-
raise ValueError(
1054-
f"Invalid input, check `node_shape` should be one of {node_shapes.keys()}"
1055-
)
1056-
if line_shape not in line_shapes:
1057-
raise ValueError(
1058-
f"Invalid input, check `line_shape` should be one of {line_shapes}"
1059-
)
1060-
if edge_arrow not in edge_arrows:
1061-
raise ValueError(
1062-
f"Invalid input, check `edge_arrow` should be one of {edge_arrows.keys()}"
1063-
)
1003+
assert_str_in_list("rankdir", rankdir, rankdirs)
1004+
assert_key_in_dict("node_shape", node_shape, node_shapes)
1005+
assert_str_in_list("line_shape", line_shape, line_shapes)
1006+
assert_key_in_dict("edge_arrow", edge_arrow, edge_arrows)
10641007

10651008
mermaid_template = """```mermaid\n{title}{line_style}\nflowchart {rankdir}\n{flows}\n{styles}\n```"""
10661009
flowchart_template = "{from_node_ref}{from_node_name}{flow_style} {arrow}{arrow_label} {to_node_ref}{to_node_name}"
@@ -1104,7 +1047,7 @@ def mermaid_name(self) -> str:
11041047
"""
11051048
if self.is_root:
11061049
return "0"
1107-
return f"{self.parent.mermaid_name}{self.parent.children.index(self)}"
1050+
return f"{self.parent.mermaid_name}-{self.parent.children.index(self)}"
11081051

11091052
tree_mermaid = clone_tree(tree, MermaidNode)
11101053
default_edge_arrow = edge_arrows[edge_arrow]

‎bigtree/utils/assertions.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from typing import Any, Dict, List
2+
3+
4+
def assert_str_in_list(
5+
parameter_name: str,
6+
parameter: Any,
7+
accepted_parameters: List[Any],
8+
) -> None:
9+
"""Raise ValueError is parameter is not in list of accepted parameters
10+
11+
Args:
12+
parameter_name (str): parameter name for error message
13+
parameter (Any): argument input for parameter
14+
accepted_parameters (List[Any]): list of accepted parameters
15+
"""
16+
if parameter not in accepted_parameters:
17+
raise ValueError(
18+
f"Invalid input, check `{parameter_name}` should be one of {accepted_parameters}"
19+
)
20+
21+
22+
def assert_key_in_dict(
23+
parameter_name: str,
24+
parameter: Any,
25+
accepted_parameters: Dict[Any, Any],
26+
) -> None:
27+
"""Raise ValueError is parameter is not in key of dictionary
28+
29+
Args:
30+
parameter_name (str): parameter name for error message
31+
parameter (Any): argument input for parameter
32+
accepted_parameters (Dict[Any]): dictionary of accepted parameters
33+
"""
34+
if parameter not in accepted_parameters:
35+
raise ValueError(
36+
f"Invalid input, check `{parameter_name}` should be one of {accepted_parameters.keys()}"
37+
)

‎bigtree/utils/constants.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
from typing import Dict, List, Tuple
2+
3+
4+
class ExportConstants:
5+
AVAILABLE_STYLES: Dict[str, Tuple[str, str, str]] = {
6+
"ansi": ("| ", "|-- ", "`-- "),
7+
"ascii": ("| ", "|-- ", "+-- "),
8+
"const": ("\u2502 ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 "),
9+
"const_bold": ("\u2503 ", "\u2523\u2501\u2501 ", "\u2517\u2501\u2501 "),
10+
"rounded": ("\u2502 ", "\u251c\u2500\u2500 ", "\u2570\u2500\u2500 "),
11+
"double": ("\u2551 ", "\u2560\u2550\u2550 ", "\u255a\u2550\u2550 "),
12+
"custom": ("", "", ""),
13+
}
14+
15+
16+
class MermaidConstants:
17+
RANK_DIR: List[str] = ["TB", "BT", "LR", "RL"]
18+
LINE_SHAPES: List[str] = [
19+
"basis",
20+
"bumpX",
21+
"bumpY",
22+
"cardinal",
23+
"catmullRom",
24+
"linear",
25+
"monotoneX",
26+
"monotoneY",
27+
"natural",
28+
"step",
29+
"stepAfter",
30+
"stepBefore",
31+
]
32+
NODE_SHAPES: Dict[str, str] = {
33+
"rounded_edge": """("{label}")""",
34+
"stadium": """(["{label}"])""",
35+
"subroutine": """[["{label}"]]""",
36+
"cylindrical": """[("{label}")]""",
37+
"circle": """(("{label}"))""",
38+
"asymmetric": """>"{label}"]""",
39+
"rhombus": """{{"{label}"}}""",
40+
"hexagon": """{{{{"{label}"}}}}""",
41+
"parallelogram": """[/"{label}"/]""",
42+
"parallelogram_alt": """[\\"{label}"\\]""",
43+
"trapezoid": """[/"{label}"\\]""",
44+
"trapezoid_alt": """[\\"{label}"/]""",
45+
"double_circle": """((("{label}")))""",
46+
}
47+
EDGE_ARROWS: Dict[str, str] = {
48+
"normal": "-->",
49+
"bold": "==>",
50+
"dotted": "-.->",
51+
"open": "---",
52+
"bold_open": "===",
53+
"dotted_open": "-.-",
54+
"invisible": "~~~",
55+
"circle": "--o",
56+
"cross": "--x",
57+
"double_normal": "<-->",
58+
"double_circle": "o--o",
59+
"double_cross": "x--x",
60+
}

0 commit comments

Comments
 (0)