Merge pull request #315 from kayjan/bugfix/mermaid-test

Fix doctest and mermaid title
kayjan · Nov 3, 2024 · 63961b2 · 63961b2
2 parents 590e1b4 + 1e184d6
commit 63961b2
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 6 deletions.
diff --git a/bigtree/tree/export.py b/bigtree/tree/export.py
@@ -1543,9 +1543,19 @@ def tree_to_mermaid(
 
         **Customize node shape, edge label, edge arrow, and custom node attributes**
 
-        >>> graph = tree_to_mermaid(root, node_shape_attr="node_shape", edge_label="edge_label", edge_arrow_attr="edge_arrow", node_attr="node_style")
+        >>> graph = tree_to_mermaid(
+        ...     root,
+        ...     title="Mermaid Diagram",
+        ...     node_shape_attr="node_shape",
+        ...     edge_label="edge_label",
+        ...     edge_arrow_attr="edge_arrow",
+        ...     node_attr="node_style",
+        ... )
         >>> print(graph)
         ```mermaid
+        ---
+        title: Mermaid Diagram
+        ---
         %%{ init: { 'flowchart': { 'curve': 'basis' } } }%%
         flowchart TB
         0{"a"} ==>|Child 1| 0-0("b")
@@ -1599,7 +1609,7 @@ def tree_to_mermaid(
     style_template = "classDef {style_name} {style}"
 
     # Content
-    title = f"---\ntitle: {title}\n---" if title else ""
+    title = f"---\ntitle: {title}\n---\n" if title else ""
     line_style = f"%%{{ init: {{ 'flowchart': {{ 'curve': '{line_shape}' }} }} }}%%"
     styles = []
     flows = []

diff --git a/bigtree/tree/helper.py b/bigtree/tree/helper.py
@@ -289,9 +289,9 @@ def get_tree_diff(
         >>> tree_diff = get_tree_diff(root, root_other)
         >>> tree_diff.show()
         Downloads
-        ├── photo2.jpg (-)
-        └── Pictures
-            └── photo2.jpg (+)
+        ├── Pictures
+        │   └── photo2.jpg (+)
+        └── photo2.jpg (-)
 
         >>> tree_diff = get_tree_diff(root, root_other, only_diff=False)
         >>> tree_diff.show()
@@ -302,7 +302,7 @@ def get_tree_diff(
         ├── file1.doc
         └── photo2.jpg (-)
 
-        >>> tree_diff = get_tree_diff(root, root_other, detail=True)
+        >>> tree_diff = get_tree_diff(root, root_other, only_diff=False, detail=True)
         >>> tree_diff.show()
         Downloads
         ├── Pictures

diff --git a/tests/tree/test_export.py b/tests/tree/test_export.py
@@ -2099,6 +2099,28 @@ def test_tree_to_mermaid(tree_node):
         )
         assert mermaid_md == expected_str
 
+    @staticmethod
+    def test_tree_to_mermaid_title(tree_node):
+        mermaid_md = export.tree_to_mermaid(tree_node, title="Mermaid Diagram")
+        expected_str = (
+            """```mermaid\n"""
+            """---\n"""
+            """title: Mermaid Diagram\n"""
+            """---\n"""
+            """%%{ init: { \'flowchart\': { \'curve\': \'basis\' } } }%%\n"""
+            """flowchart TB\n"""
+            """0("a") --> 0-0("b")\n"""
+            """0-0 --> 0-0-0("d")\n"""
+            """0-0 --> 0-0-1("e")\n"""
+            """0-0-1 --> 0-0-1-0("g")\n"""
+            """0-0-1 --> 0-0-1-1("h")\n"""
+            """0("a") --> 0-1("c")\n"""
+            """0-1 --> 0-1-0("f")\n"""
+            """classDef default stroke-width:1\n"""
+            """```"""
+        )
+        assert mermaid_md == expected_str
+
     @staticmethod
     def test_tree_to_mermaid_invalid_rankdir_error(tree_node):
         with pytest.raises(ValueError) as exc_info: