add node_shape to tree_to_dot method + corresponding test case

CHANGELOG.md

@@ -5,6 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [0.7.2] - 2023-02-18
+### Added
+- Tree Export: Added `node_shape` parameter in `tree_to_dot` export function for easier way to customize node shape.
+
 ## [0.7.1] - 2023-02-18
 ### Added
 - BaseNode/Node: Added `go_to` BaseNode method to travel from one node to another node from the same tree.
@@ -191,6 +195,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Utility Iterator: Tree traversal methods.
 - Workflow To Do App: Tree use case with to-do list implementation.
 
+[0.7.2]: https://github.com/kayjan/bigtree/compare/v0.7.1...v0.7.2
 [0.7.1]: https://github.com/kayjan/bigtree/compare/v0.7.0...v0.7.1
 [0.7.0]: https://github.com/kayjan/bigtree/compare/v0.6.10...v0.7.0
 [0.6.10]: https://github.com/kayjan/bigtree/compare/v0.6.9...v0.6.10

bigtree/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.7.1"
+__version__ = "0.7.2"
 
 from bigtree.binarytree.construct import list_to_binarytree
 from bigtree.dag.construct import dataframe_to_dag, dict_to_dag, list_to_dag

bigtree/tree/export.py

@@ -596,6 +596,7 @@ def tree_to_dot(
     rankdir: str = "TB",
     bg_colour: str = None,
     node_colour: str = None,
+    node_shape: str = None,
     edge_colour: str = None,
     node_attr: str = None,
     edge_attr: str = None,
@@ -631,19 +632,23 @@ def tree_to_dot(
     ...
     ...     @property
     ...     def edge_attr(self):
-    ...         return {"label": self.edge_label}
+    ...         if self.edge_label:
+    ...             return {"label": self.edge_label}
+    ...         return {}
     ...
     ...     @property
     ...     def node_attr(self):
-    ...         return {"shape": self.node_shape}
+    ...         if self.node_shape:
+    ...             return {"shape": self.node_shape}
+    ...         return {}
     >>>
     >>>
     >>> root = CustomNode("a", node_shape="circle")
-    >>> b = CustomNode("b", node_shape="diamond", edge_label="child", parent=root)
-    >>> c = CustomNode("c", node_shape="diamond", edge_label="child", parent=root)
+    >>> b = CustomNode("b", edge_label="child", parent=root)
+    >>> c = CustomNode("c", edge_label="child", parent=root)
     >>> d = CustomNode("d", node_shape="square", edge_label="child", parent=b)
     >>> e = CustomNode("e", node_shape="square", edge_label="child", parent=b)
-    >>> graph = tree_to_dot(root, node_colour="gold", node_attr="node_attr", edge_attr="edge_attr")
+    >>> graph = tree_to_dot(root, node_colour="gold", node_shape="diamond", node_attr="node_attr", edge_attr="edge_attr")
     >>> graph.write_png("assets/custom_tree.png")
 
     .. image:: https://github.com/kayjan/bigtree/raw/master/assets/custom_tree.png
@@ -655,10 +660,12 @@ def tree_to_dot(
             'LR' (left to right), 'RL' (right to left)
         bg_colour (str): background color of image, defaults to None
         node_colour (str): fill colour of nodes, defaults to None
+        node_shape (str): shape of nodes, defaults to None
+            Possible node_shape include "circle", "square", "diamond", "triangle"
         edge_colour (str): colour of edges, defaults to None
-        node_attr (str): `Node` attribute for node style, overrides node_colour, defaults to None.
+        node_attr (str): `Node` attribute for node style, overrides `node_colour` and `node_shape`, defaults to None.
             Possible node style (attribute value) include {"style": "filled", "fillcolor": "gold", "shape": "diamond"}
-        edge_attr (str): `Node` attribute for edge style, overrides edge_colour, defaults to None.
+        edge_attr (str): `Node` attribute for edge style, overrides `edge_colour`, defaults to None.
             Possible edge style (attribute value) include {"style": "bold", "label": "edge label", "color": "black"}
 
     Returns:
@@ -682,6 +689,9 @@ def tree_to_dot(
     else:
         node_style = dict()
 
+    if node_shape:
+        node_style["shape"] = node_shape
+
     if edge_colour:
         edge_style = dict(color=edge_colour)
     else:

tests/tree/test_export.py

@@ -805,7 +805,7 @@ def test_tree_to_dot_bg_color(tree_node):
             ), f"Expected {expected_str} not in actual string"
 
     @staticmethod
-    def test_tree_to_dot_fill_color(tree_node):
+    def test_tree_to_dot_node_colour(tree_node):
         graph = tree_to_dot(tree_node, node_colour="gold")
         expected = """strict digraph G {\nrankdir=TB;\na0 [fillcolor=gold, label=a, style=filled];\nb0 [fillcolor=gold, label=b, style=filled];\na0 -> b0;\nd0 [fillcolor=gold, label=d, style=filled];\nb0 -> d0;\ne0 [fillcolor=gold, label=e, style=filled];\nb0 -> e0;\ng0 [fillcolor=gold, label=g, style=filled];\ne0 -> g0;\nh0 [fillcolor=gold, label=h, style=filled];\ne0 -> h0;\nc0 [fillcolor=gold, label=c, style=filled];\na0 -> c0;\nf0 [fillcolor=gold, label=f, style=filled];\nc0 -> f0;\n}\n"""
         actual = graph.to_string()
@@ -815,6 +815,17 @@ def test_tree_to_dot_fill_color(tree_node):
                 expected_str in actual
             ), f"Expected {expected_str} not in actual string"
 
+    @staticmethod
+    def test_tree_to_dot_node_shape(tree_node):
+        graph = tree_to_dot(tree_node, node_shape="triangle")
+        expected = """strict digraph G {\nrankdir=TB;\na0 [label=a, shape=triangle];\nb0 [label=b, shape=triangle];\na0 -> b0;\nd0 [label=d, shape=triangle];\nb0 -> d0;\ne0 [label=e, shape=triangle];\nb0 -> e0;\ng0 [label=g, shape=triangle];\ne0 -> g0;\nh0 [label=h, shape=triangle];\ne0 -> h0;\nc0 [label=c, shape=triangle];\na0 -> c0;\nf0 [label=f, shape=triangle];\nc0 -> f0;\n}\n"""
+        actual = graph.to_string()
+        graph.write_png("tests/tree_triangle.png")
+        for expected_str in expected.split():
+            assert (
+                expected_str in actual
+            ), f"Expected {expected_str} not in actual string"
+
     @staticmethod
     def test_tree_to_dot_edge_colour(tree_node):
         graph = tree_to_dot(tree_node, edge_colour="red")