Change: Mermaid delimiter, abstract assertions and constants to anoth…

README.md

@@ -915,10 +915,10 @@ print(mermaid_md)
 ```mermaid
 %%{ init: { 'flowchart': { 'curve': 'basis' } } }%%
 flowchart TB
-0("a") --> 00("b")
-00 --> 000("d")
-00 --> 001("e")
-0("a") --> 01("c")
+0("a") --> 0-0("b")
+0-0 --> 0-0-0("d")
+0-0 --> 0-0-1("e")
+0("a") --> 0-1("c")
 classDef default stroke-width:1
 ```
 

bigtree/tree/export.py

@@ -6,6 +6,8 @@
 
 from bigtree.node.node import Node
 from bigtree.tree.search import find_path
+from bigtree.utils.assertions import assert_key_in_dict, assert_str_in_list
+from bigtree.utils.constants import ExportConstants, MermaidConstants
 from bigtree.utils.exceptions import (
     optional_dependencies_image,
     optional_dependencies_pandas,
@@ -41,16 +43,6 @@
 
 T = TypeVar("T", bound=Node)
 
-available_styles = {
-    "ansi": ("|   ", "|-- ", "`-- "),
-    "ascii": ("|   ", "|-- ", "+-- "),
-    "const": ("\u2502   ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 "),
-    "const_bold": ("\u2503   ", "\u2523\u2501\u2501 ", "\u2517\u2501\u2501 "),
-    "rounded": ("\u2502   ", "\u251c\u2500\u2500 ", "\u2570\u2500\u2500 "),
-    "double": ("\u2551   ", "\u2560\u2550\u2550 ", "\u255a\u2550\u2550 "),
-    "custom": ("", "", ""),
-}
-
 
 def print_tree(
     tree: T,
@@ -324,6 +316,7 @@ def yield_tree(
         style (str): style of print, defaults to abstract style
         custom_style (Iterable[str]): style of stem, branch and final stem, used when `style` is set to 'custom'
     """
+    available_styles = ExportConstants.AVAILABLE_STYLES
     if style not in available_styles.keys():
         raise ValueError(
             f"Choose one of {available_styles.keys()} style, use `custom` to define own style"
@@ -1001,66 +994,16 @@ def tree_to_mermaid(
     """
     from bigtree.tree.helper import clone_tree
 
-    rankdirs = ["TB", "BT", "LR", "RL"]
-    line_shapes = [
-        "basis",
-        "bumpX",
-        "bumpY",
-        "cardinal",
-        "catmullRom",
-        "linear",
-        "monotoneX",
-        "monotoneY",
-        "natural",
-        "step",
-        "stepAfter",
-        "stepBefore",
-    ]
-    node_shapes = {
-        "rounded_edge": """("{label}")""",
-        "stadium": """(["{label}"])""",
-        "subroutine": """[["{label}"]]""",
-        "cylindrical": """[("{label}")]""",
-        "circle": """(("{label}"))""",
-        "asymmetric": """>"{label}"]""",
-        "rhombus": """{{"{label}"}}""",
-        "hexagon": """{{{{"{label}"}}}}""",
-        "parallelogram": """[/"{label}"/]""",
-        "parallelogram_alt": """[\\"{label}"\\]""",
-        "trapezoid": """[/"{label}"\\]""",
-        "trapezoid_alt": """[\\"{label}"/]""",
-        "double_circle": """((("{label}")))""",
-    }
-    edge_arrows = {
-        "normal": "-->",
-        "bold": "==>",
-        "dotted": "-.->",
-        "open": "---",
-        "bold_open": "===",
-        "dotted_open": "-.-",
-        "invisible": "~~~",
-        "circle": "--o",
-        "cross": "--x",
-        "double_normal": "<-->",
-        "double_circle": "o--o",
-        "double_cross": "x--x",
-    }
+    rankdirs = MermaidConstants.RANK_DIR
+    line_shapes = MermaidConstants.LINE_SHAPES
+    node_shapes = MermaidConstants.NODE_SHAPES
+    edge_arrows = MermaidConstants.EDGE_ARROWS
 
     # Assertions
-    if rankdir not in rankdirs:
-        raise ValueError(f"Invalid input, check `rankdir` should be one of {rankdirs}")
-    if node_shape not in node_shapes:
-        raise ValueError(
-            f"Invalid input, check `node_shape` should be one of {node_shapes.keys()}"
-        )
-    if line_shape not in line_shapes:
-        raise ValueError(
-            f"Invalid input, check `line_shape` should be one of {line_shapes}"
-        )
-    if edge_arrow not in edge_arrows:
-        raise ValueError(
-            f"Invalid input, check `edge_arrow` should be one of {edge_arrows.keys()}"
-        )
+    assert_str_in_list("rankdir", rankdir, rankdirs)
+    assert_key_in_dict("node_shape", node_shape, node_shapes)
+    assert_str_in_list("line_shape", line_shape, line_shapes)
+    assert_key_in_dict("edge_arrow", edge_arrow, edge_arrows)
 
     mermaid_template = """```mermaid\n{title}{line_style}\nflowchart {rankdir}\n{flows}\n{styles}\n```"""
     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:
             """
             if self.is_root:
                 return "0"
-            return f"{self.parent.mermaid_name}{self.parent.children.index(self)}"
+            return f"{self.parent.mermaid_name}-{self.parent.children.index(self)}"
 
     tree_mermaid = clone_tree(tree, MermaidNode)
     default_edge_arrow = edge_arrows[edge_arrow]

bigtree/utils/assertions.py

@@ -0,0 +1,37 @@
+from typing import Any, Dict, List
+
+
+def assert_str_in_list(
+    parameter_name: str,
+    parameter: Any,
+    accepted_parameters: List[Any],
+) -> None:
+    """Raise ValueError is parameter is not in list of accepted parameters
+
+    Args:
+        parameter_name (str): parameter name for error message
+        parameter (Any): argument input for parameter
+        accepted_parameters (List[Any]): list of accepted parameters
+    """
+    if parameter not in accepted_parameters:
+        raise ValueError(
+            f"Invalid input, check `{parameter_name}` should be one of {accepted_parameters}"
+        )
+
+
+def assert_key_in_dict(
+    parameter_name: str,
+    parameter: Any,
+    accepted_parameters: Dict[Any, Any],
+) -> None:
+    """Raise ValueError is parameter is not in key of dictionary
+
+    Args:
+        parameter_name (str): parameter name for error message
+        parameter (Any): argument input for parameter
+        accepted_parameters (Dict[Any]): dictionary of accepted parameters
+    """
+    if parameter not in accepted_parameters:
+        raise ValueError(
+            f"Invalid input, check `{parameter_name}` should be one of {accepted_parameters.keys()}"
+        )

bigtree/utils/constants.py

@@ -0,0 +1,60 @@
+from typing import Dict, List, Tuple
+
+
+class ExportConstants:
+    AVAILABLE_STYLES: Dict[str, Tuple[str, str, str]] = {
+        "ansi": ("|   ", "|-- ", "`-- "),
+        "ascii": ("|   ", "|-- ", "+-- "),
+        "const": ("\u2502   ", "\u251c\u2500\u2500 ", "\u2514\u2500\u2500 "),
+        "const_bold": ("\u2503   ", "\u2523\u2501\u2501 ", "\u2517\u2501\u2501 "),
+        "rounded": ("\u2502   ", "\u251c\u2500\u2500 ", "\u2570\u2500\u2500 "),
+        "double": ("\u2551   ", "\u2560\u2550\u2550 ", "\u255a\u2550\u2550 "),
+        "custom": ("", "", ""),
+    }
+
+
+class MermaidConstants:
+    RANK_DIR: List[str] = ["TB", "BT", "LR", "RL"]
+    LINE_SHAPES: List[str] = [
+        "basis",
+        "bumpX",
+        "bumpY",
+        "cardinal",
+        "catmullRom",
+        "linear",
+        "monotoneX",
+        "monotoneY",
+        "natural",
+        "step",
+        "stepAfter",
+        "stepBefore",
+    ]
+    NODE_SHAPES: Dict[str, str] = {
+        "rounded_edge": """("{label}")""",
+        "stadium": """(["{label}"])""",
+        "subroutine": """[["{label}"]]""",
+        "cylindrical": """[("{label}")]""",
+        "circle": """(("{label}"))""",
+        "asymmetric": """>"{label}"]""",
+        "rhombus": """{{"{label}"}}""",
+        "hexagon": """{{{{"{label}"}}}}""",
+        "parallelogram": """[/"{label}"/]""",
+        "parallelogram_alt": """[\\"{label}"\\]""",
+        "trapezoid": """[/"{label}"\\]""",
+        "trapezoid_alt": """[\\"{label}"/]""",
+        "double_circle": """((("{label}")))""",
+    }
+    EDGE_ARROWS: Dict[str, str] = {
+        "normal": "-->",
+        "bold": "==>",
+        "dotted": "-.->",
+        "open": "---",
+        "bold_open": "===",
+        "dotted_open": "-.-",
+        "invisible": "~~~",
+        "circle": "--o",
+        "cross": "--x",
+        "double_normal": "<-->",
+        "double_circle": "o--o",
+        "double_cross": "x--x",
+    }