test: Add polars test for binary tree

Author: kayjan

tests/binarytree/test_construct.py

@@ -1,6 +1,7 @@
 import unittest
 
 import pandas as pd
+import polars as pl
 import pytest
 
 from bigtree.binarytree.construct import list_to_binarytree
@@ -12,6 +13,8 @@
     list_to_tree,
     list_to_tree_by_relation,
     nested_dict_to_tree,
+    polars_to_tree,
+    polars_to_tree_by_relation,
 )
 from tests.node.test_binarynode import assert_binarytree_structure_root2
 from tests.test_constants import Constants
@@ -234,3 +237,62 @@ def tearDown(self):
     def test_dataframe_to_tree_by_relation(self):
         root = dataframe_to_tree_by_relation(self.relation_data, node_type=BinaryNode)
         assert_binarytree_structure_root2(root)
+
+
+class TestPolarsToTree(unittest.TestCase):
+    def setUp(self):
+        """
+        Binary Tree should have structure
+        1
+        ├── 2
+        │   ├── 4
+        │   │   └── 8
+        │   └── 5
+        └── 3
+            ├── 6
+            └── 7
+        """
+        self.path_data = pl.DataFrame(
+            [
+                ["1", 90],
+                ["1/2", 65],
+                ["1/3", 60],
+                ["1/2/4", 40],
+                ["1/2/5", 35],
+                ["1/3/6", 38],
+                ["1/3/7", 10],
+                ["1/2/4/8", 6],
+            ],
+            schema=["PATH", "age"],
+        )
+
+    def tearDown(self):
+        self.path_data = None
+
+    def test_polars_to_tree(self):
+        root = polars_to_tree(self.path_data, node_type=BinaryNode)
+        assert_binarytree_structure_root2(root)
+
+
+class TestPolarsToTreeByRelation(unittest.TestCase):
+    def setUp(self):
+        self.relation_data = pl.DataFrame(
+            [
+                ["1", None, 90],
+                ["2", "1", 65],
+                ["3", "1", 60],
+                ["4", "2", 40],
+                ["5", "2", 35],
+                ["6", "3", 38],
+                ["7", "3", 10],
+                ["8", "4", 6],
+            ],
+            schema=["child", "parent", "age"],
+        )
+
+    def tearDown(self):
+        self.relation_data = None
+
+    def test_polars_to_tree_by_relation(self):
+        root = polars_to_tree_by_relation(self.relation_data, node_type=BinaryNode)
+        assert_binarytree_structure_root2(root)

tests/binarytree/test_export.py

@@ -1,4 +1,5 @@
 import pandas as pd
+import polars as pl
 
 from bigtree.tree.export import (
     print_tree,
@@ -7,6 +8,7 @@
     tree_to_dot,
     tree_to_nested_dict,
     tree_to_newick,
+    tree_to_polars,
 )
 from tests.conftest import assert_print_statement
 from tests.test_constants import Constants
@@ -41,6 +43,26 @@ def test_tree_to_dataframe(binarytree_node):
         pd.testing.assert_frame_equal(expected, actual)
 
 
+class TestTreeToPolars:
+    @staticmethod
+    def test_tree_to_polars(binarytree_node):
+        expected = pl.DataFrame(
+            [
+                ["/1", "1"],
+                ["/1/2", "2"],
+                ["/1/2/4", "4"],
+                ["/1/2/4/8", "8"],
+                ["/1/2/5", "5"],
+                ["/1/3", "3"],
+                ["/1/3/6", "6"],
+                ["/1/3/7", "7"],
+            ],
+            schema=["path", "name"],
+        )
+        actual = tree_to_polars(binarytree_node)
+        assert expected.equals(actual)
+
+
 class TestTreeToDict:
     @staticmethod
     def test_tree_to_dict(binarytree_node):