Skip to content

Commit 1e184d6

Browse files
committed
fix: doctest and mermaid title
1 parent 590e1b4 commit 1e184d6

File tree

3 files changed

+38
-6
lines changed

3 files changed

+38
-6
lines changed

bigtree/tree/export.py

+12-2
Original file line numberDiff line numberDiff line change
@@ -1543,9 +1543,19 @@ def tree_to_mermaid(
15431543
15441544
**Customize node shape, edge label, edge arrow, and custom node attributes**
15451545
1546-
>>> graph = tree_to_mermaid(root, node_shape_attr="node_shape", edge_label="edge_label", edge_arrow_attr="edge_arrow", node_attr="node_style")
1546+
>>> graph = tree_to_mermaid(
1547+
... root,
1548+
... title="Mermaid Diagram",
1549+
... node_shape_attr="node_shape",
1550+
... edge_label="edge_label",
1551+
... edge_arrow_attr="edge_arrow",
1552+
... node_attr="node_style",
1553+
... )
15471554
>>> print(graph)
15481555
```mermaid
1556+
---
1557+
title: Mermaid Diagram
1558+
---
15491559
%%{ init: { 'flowchart': { 'curve': 'basis' } } }%%
15501560
flowchart TB
15511561
0{"a"} ==>|Child 1| 0-0("b")
@@ -1599,7 +1609,7 @@ def tree_to_mermaid(
15991609
style_template = "classDef {style_name} {style}"
16001610

16011611
# Content
1602-
title = f"---\ntitle: {title}\n---" if title else ""
1612+
title = f"---\ntitle: {title}\n---\n" if title else ""
16031613
line_style = f"%%{{ init: {{ 'flowchart': {{ 'curve': '{line_shape}' }} }} }}%%"
16041614
styles = []
16051615
flows = []

bigtree/tree/helper.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -289,9 +289,9 @@ def get_tree_diff(
289289
>>> tree_diff = get_tree_diff(root, root_other)
290290
>>> tree_diff.show()
291291
Downloads
292-
├── photo2.jpg (-)
293-
└── Pictures
294-
└── photo2.jpg (+)
292+
├── Pictures
293+
└── photo2.jpg (+)
294+
└── photo2.jpg (-)
295295
296296
>>> tree_diff = get_tree_diff(root, root_other, only_diff=False)
297297
>>> tree_diff.show()
@@ -302,7 +302,7 @@ def get_tree_diff(
302302
├── file1.doc
303303
└── photo2.jpg (-)
304304
305-
>>> tree_diff = get_tree_diff(root, root_other, detail=True)
305+
>>> tree_diff = get_tree_diff(root, root_other, only_diff=False, detail=True)
306306
>>> tree_diff.show()
307307
Downloads
308308
├── Pictures

tests/tree/test_export.py

+22
Original file line numberDiff line numberDiff line change
@@ -2099,6 +2099,28 @@ def test_tree_to_mermaid(tree_node):
20992099
)
21002100
assert mermaid_md == expected_str
21012101

2102+
@staticmethod
2103+
def test_tree_to_mermaid_title(tree_node):
2104+
mermaid_md = export.tree_to_mermaid(tree_node, title="Mermaid Diagram")
2105+
expected_str = (
2106+
"""```mermaid\n"""
2107+
"""---\n"""
2108+
"""title: Mermaid Diagram\n"""
2109+
"""---\n"""
2110+
"""%%{ init: { \'flowchart\': { \'curve\': \'basis\' } } }%%\n"""
2111+
"""flowchart TB\n"""
2112+
"""0("a") --> 0-0("b")\n"""
2113+
"""0-0 --> 0-0-0("d")\n"""
2114+
"""0-0 --> 0-0-1("e")\n"""
2115+
"""0-0-1 --> 0-0-1-0("g")\n"""
2116+
"""0-0-1 --> 0-0-1-1("h")\n"""
2117+
"""0("a") --> 0-1("c")\n"""
2118+
"""0-1 --> 0-1-0("f")\n"""
2119+
"""classDef default stroke-width:1\n"""
2120+
"""```"""
2121+
)
2122+
assert mermaid_md == expected_str
2123+
21022124
@staticmethod
21032125
def test_tree_to_mermaid_invalid_rankdir_error(tree_node):
21042126
with pytest.raises(ValueError) as exc_info:

0 commit comments

Comments
 (0)