Clean up (#399)

* feat: construct trees with nested_dict_key_to_tree

* refactor: test clean up

* feat: export trees with tree_to_nested_dict_key

* feat: add docs

* docs: update CHANGELOG

* docs: enhance docstring

* refactor: abstract out export tree logic

* docs: update CHANGELOG

* refactor: shift to common file

* refactor: shift to common file

* refactor: shift to common file

* docs: misc fix

* refactor: clean up

* refactor: clean up + error check for dataframe_to_dag

* refactor: clean up

* refactor: enhance assemble_attributes

* refactor: enhance assemble_attributes

* test: fix test coverage

Author: kayjan

CHANGELOG.md

@@ -6,7 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Changed
-- Misc: Some code refactoring
+- Misc: Some code refactoring, enhance assemble_attributes.
+### Fixed
+- Test: Test for `tree_to_nested_dict_key` for BinaryNode.
+- Error: Check and throw error for `dataframe_to_dag`, previously this error will not have been found out.
 
 ## [0.30.0] - 2025-09-05
 ### Added:

bigtree/dag/construct.py

@@ -40,23 +40,14 @@ def list_to_dag(
     assertions.assert_length_not_empty(relations, "Input list", "relations")
 
     node_dict: Dict[str, T] = dict()
-    parent_node: T = dagnode.DAGNode()  # type: ignore[assignment]
+    child_name: str = ""
 
     for parent_name, child_name in relations:
-        if parent_name not in node_dict:
-            parent_node = node_type(parent_name)
-            node_dict[parent_name] = parent_node
-        else:
-            parent_node = node_dict[parent_name]
-        if child_name not in node_dict:
-            child_node = node_type(child_name)
-            node_dict[child_name] = child_node
-        else:
-            child_node = node_dict[child_name]
+        node_dict[parent_name] = node_dict.get(parent_name, node_type(parent_name))
+        node_dict[child_name] = node_dict.get(child_name, node_type(child_name))
+        node_dict[child_name].parents = [node_dict[parent_name]]
 
-        child_node.parents = [parent_node]
-
-    return parent_node
+    return node_dict[child_name]
 
 
 def dict_to_dag(
@@ -92,35 +83,27 @@ def dict_to_dag(
     assertions.assert_length_not_empty(relation_attrs, "Dictionary", "relation_attrs")
 
     node_dict: Dict[str, T] = dict()
-    parent_node: T | None = None
+    _parent_name: Optional[str] = None
 
     for child_name, node_attrs in relation_attrs.items():
         node_attrs = node_attrs.copy()
-        parent_names: List[str] = []
-        if parent_key in node_attrs:
-            parent_names = node_attrs.pop(parent_key)
+        parent_names = node_attrs.pop(parent_key, [])
         assertions.assert_not_reserved_keywords(
             node_attrs, ["parent", "parents", "children"]
         )
 
-        if child_name in node_dict:
-            child_node = node_dict[child_name]
-            child_node.set_attrs(node_attrs)
-        else:
-            child_node = node_type(child_name, **node_attrs)
-            node_dict[child_name] = child_node
+        node_dict[child_name] = node_dict.get(child_name, node_type(child_name))
+        node_dict[child_name].set_attrs(node_attrs)
 
         for parent_name in parent_names:
-            parent_node = node_dict.get(parent_name, node_type(parent_name))
-            node_dict[parent_name] = parent_node
-            child_node.parents = [parent_node]
+            node_dict[parent_name] = node_dict.get(parent_name, node_type(parent_name))
+            node_dict[child_name].parents = [node_dict[parent_name]]
+            _parent_name = parent_name
 
-    if parent_node is None:
-        raise ValueError(
-            f"Parent key {parent_key} not in dictionary, check `relation_attrs` and `parent_key`"
-        )
+    if _parent_name is None:
+        raise ValueError("No parent specified, check `relation_attrs` and `parent_key`")
 
-    return parent_node
+    return node_dict[_parent_name]
 
 
 @exceptions.optional_dependencies_pandas
@@ -196,21 +179,23 @@ def dataframe_to_dag(
         raise ValueError(f"Child name cannot be empty, check column: {child_col}")
 
     node_dict: Dict[str, T] = dict()
-    parent_node: T = dagnode.DAGNode()  # type: ignore[assignment]
+    _parent_name: Optional[str] = None
 
     for row in data.reset_index(drop=True).to_dict(orient="index").values():
         child_name = row[child_col]
         parent_name = row[parent_col]
         node_attrs = common.filter_attributes(
             row, omit_keys=["name", child_col, parent_col], omit_null_values=True
         )
-        child_node = node_dict.get(child_name, node_type(child_name, **node_attrs))
-        child_node.set_attrs(node_attrs)
-        node_dict[child_name] = child_node
+        node_dict[child_name] = node_dict.get(child_name, node_type(child_name))
+        node_dict[child_name].set_attrs(node_attrs)
 
         if not common.isnull(parent_name):
-            parent_node = node_dict.get(parent_name, node_type(parent_name))
-            node_dict[parent_name] = parent_node
-            child_node.parents = [parent_node]
+            node_dict[parent_name] = node_dict.get(parent_name, node_type(parent_name))
+            node_dict[child_name].parents = [node_dict[parent_name]]
+            _parent_name = parent_name
+
+    if _parent_name is None:
+        raise ValueError("No parent specified, check `data` and `parent_col`")
 
-    return parent_node
+    return node_dict[_parent_name]

bigtree/dag/export.py

@@ -85,18 +85,17 @@ def dag_to_dict(
 
     for parent_node, child_node in iterators.dag_iterator(dag):
         if parent_node.is_root:
-            data_parent: Dict[str, Any] = {}
-            data_parent = common.assemble_attributes(
-                parent_node, attr_dict, all_attrs, data_parent
-            )
+            data_parent = common.assemble_attributes(parent_node, attr_dict, all_attrs)
             data_dict[parent_node.node_name] = data_parent
 
         if data_dict.get(child_node.node_name):
             data_dict[child_node.node_name][parent_key].append(parent_node.node_name)
         else:
-            data_child = {parent_key: [parent_node.node_name]}
             data_child = common.assemble_attributes(
-                child_node, attr_dict, all_attrs, data_child
+                child_node,
+                attr_dict,
+                all_attrs,
+                parent_col=(parent_key, [parent_node.node_name]),
             )
             data_dict[child_node.node_name] = data_child
     return data_dict
@@ -144,15 +143,21 @@ def dag_to_dataframe(
 
     for parent_node, child_node in iterators.dag_iterator(dag):
         if parent_node.is_root:
-            data_parent = {name_col: parent_node.node_name, parent_col: None}
             data_parent = common.assemble_attributes(
-                parent_node, attr_dict, all_attrs, data_parent
+                parent_node,
+                attr_dict,
+                all_attrs,
+                name_col=name_col,
+                parent_col=(parent_col, None),
             )
             data_list.append(data_parent)
 
-        data_child = {name_col: child_node.node_name, parent_col: parent_node.node_name}
         data_child = common.assemble_attributes(
-            child_node, attr_dict, all_attrs, data_child
+            child_node,
+            attr_dict,
+            all_attrs,
+            name_col=name_col,
+            parent_col=(parent_col, parent_node.node_name),
         )
         data_list.append(data_child)
     return pd.DataFrame(data_list).drop_duplicates().reset_index(drop=True)

bigtree/tree/construct/dataframes.py

@@ -576,8 +576,8 @@ def _recursive_add_child(parent_node: T) -> None:
         """
         child_rows = data[data[parent_col] == parent_node.node_name]
 
-        for row in child_rows.to_dict(orient="index").values():
-            child_node = node_type(**_retrieve_attr(row))
+        for _row in child_rows.to_dict(orient="index").values():
+            child_node = node_type(**_retrieve_attr(_row))
             child_node.parent = parent_node
             _recursive_add_child(child_node)
 

bigtree/tree/export/dataframes.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, TypeVar
+from typing import Dict, Optional, TypeVar
 
 from bigtree.node import node
 from bigtree.utils import common, exceptions
@@ -95,19 +95,13 @@ def _recursive_append(_node: T) -> None:
                 and (not skip_depth or _node.depth > skip_depth)
                 and (not leaf_only or _node.is_leaf)
             ):
-                data_child: Dict[str, Any] = {}
-                if path_col:
-                    data_child[path_col] = _node.path_name
-                if name_col:
-                    data_child[name_col] = _node.node_name
-                if parent_col:
-                    parent_name = None
-                    if _node.parent:
-                        parent_name = _node.parent.node_name
-                    data_child[parent_col] = parent_name
-
                 data_child = common.assemble_attributes(
-                    _node, attr_dict, all_attrs, data_child
+                    _node,
+                    attr_dict,
+                    all_attrs,
+                    path_col=path_col,
+                    name_col=name_col,
+                    parent_col=parent_col,
                 )
                 data_list.append(data_child)
             for _child in _node.children:
@@ -196,19 +190,13 @@ def _recursive_append(_node: T) -> None:
                 and (not skip_depth or _node.depth > skip_depth)
                 and (not leaf_only or _node.is_leaf)
             ):
-                data_child: Dict[str, Any] = {}
-                if path_col:
-                    data_child[path_col] = _node.path_name
-                if name_col:
-                    data_child[name_col] = _node.node_name
-                if parent_col:
-                    parent_name = None
-                    if _node.parent:
-                        parent_name = _node.parent.node_name
-                    data_child[parent_col] = parent_name
-
                 data_child = common.assemble_attributes(
-                    _node, attr_dict, all_attrs, data_child
+                    _node,
+                    attr_dict,
+                    all_attrs,
+                    path_col=path_col,
+                    name_col=name_col,
+                    parent_col=parent_col,
                 )
                 data_list.append(data_child)
             for _child in _node.children:

bigtree/tree/export/dictionaries.py

@@ -14,34 +14,6 @@
 T = TypeVar("T", bound=node.Node)
 
 
-def _assemble_attributes(
-    _node: T,
-    attr_dict: Optional[Dict[str, str]],
-    all_attrs: bool,
-    data_child: Dict[str, Any] = None,
-) -> Dict[str, Any]:
-    """Assemble attributes of node into a dictionary.
-
-    Args:
-        _node: node
-        attr_dict: node attributes mapped to dictionary key, key: node attributes, value: corresponding dictionary key
-        all_attrs: indicator whether to retrieve all ``Node`` attributes, overrides `attr_dict`
-        data_child: existing attributes, if any
-
-    Returns:
-        node attributes
-    """
-    data_child = data_child or {}
-    if all_attrs:
-        data_child.update(
-            dict(_node.describe(exclude_attributes=["name"], exclude_prefix="_"))
-        )
-    elif attr_dict:
-        for k, v in attr_dict.items():
-            data_child[v] = _node.get_attr(k)
-    return data_child
-
-
 def tree_to_dict(
     tree: T,
     name_key: Optional[str] = "name",
@@ -100,16 +72,12 @@ def _recursive_append(_node: T) -> None:
                 and (not skip_depth or _node.depth > skip_depth)
                 and (not leaf_only or _node.is_leaf)
             ):
-                data_child: Dict[str, Any] = {}
-                if name_key:
-                    data_child[name_key] = _node.node_name
-                if parent_key:
-                    parent_name = None
-                    if _node.parent:
-                        parent_name = _node.parent.node_name
-                    data_child[parent_key] = parent_name
                 data_child = common.assemble_attributes(
-                    _node, attr_dict, all_attrs, data_child
+                    _node,
+                    attr_dict,
+                    all_attrs,
+                    name_col=name_key,
+                    parent_col=parent_key,
                 )
                 data_dict[_node.path_name] = data_child
             for _child in _node.children:
@@ -165,9 +133,8 @@ def _recursive_append(_node: T, parent_dict: Dict[str, Any]) -> None:
         """
         if _node:
             if not max_depth or _node.depth <= max_depth:
-                data_child = {name_key: _node.node_name}
                 data_child = common.assemble_attributes(
-                    _node, attr_dict, all_attrs, data_child
+                    _node, attr_dict, all_attrs, name_col=name_key
                 )
                 if child_key in parent_dict:
                     parent_dict[child_key].append(data_child)

bigtree/utils/common.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Collection, Dict, Mapping, Optional, TypeVar, Union
+from typing import Any, Collection, Dict, Mapping, Optional, Tuple, TypeVar, Union
 
 from bigtree.node import dagnode, node
 
@@ -57,25 +57,50 @@ def assemble_attributes(
     _node: T,
     attr_dict: Optional[Mapping[str, str]],
     all_attrs: bool,
-    existing_data: Dict[str, Any] = None,
+    path_col: Optional[str] = None,
+    name_col: Optional[str] = None,
+    parent_col: Optional[Union[str, Tuple[str, Any]]] = None,
 ) -> Dict[str, Any]:
     """Assemble attributes of node into a dictionary.
 
     Args:
         _node: node
         attr_dict: node attributes mapped to dictionary key, key: node attributes, value: corresponding dictionary key
         all_attrs: indicator whether to retrieve all ``Node`` attributes, overrides `attr_dict`
-        existing_data: existing attributes, if any
+        path_col: column name for `_node.path_name`, if present
+        name_col: column name for `_node.node_name`, if present
+        parent_col: if Node, column name for `_node.parent.node_name`. If DAGNode, tuple of column name and value for
+            `_node.parent.node_name`.
 
     Returns:
         node attributes
     """
-    data_attrs = existing_data or {}
+    data_attrs = {}
+
+    # Main attributes
+    if path_col:
+        assert isinstance(_node, node.Node)
+        data_attrs[path_col] = _node.path_name
+    if name_col:
+        data_attrs[name_col] = _node.node_name
+    if parent_col:
+        if isinstance(_node, node.Node):
+            assert isinstance(parent_col, str)
+            parent_name = None
+            if _node.parent:
+                parent_name = _node.parent.node_name
+            data_attrs[parent_col] = parent_name
+        else:
+            assert isinstance(parent_col, tuple)
+            data_attrs[parent_col[0]] = parent_col[1]
+
+    # Other attributes
     if all_attrs:
         data_attrs.update(
             dict(_node.describe(exclude_attributes=["name"], exclude_prefix="_"))
         )
     elif attr_dict:
         for k, v in attr_dict.items():
             data_attrs[v] = _node.get_attr(k)
+
     return data_attrs

tests/binarytree/test_export.py

@@ -92,6 +92,26 @@ def test_tree_to_nested_dict(binarytree_node):
         assert actual == expected, f"Expected\n{expected}\nReceived\n{actual}"
 
 
+class TestTreeToNestedDictKey:
+    @staticmethod
+    def test_tree_to_nested_dict_key(binarytree_node):
+        expected = {
+            "1": {
+                "children": {
+                    "2": {
+                        "children": {
+                            "4": {"children": {"8": {}}},
+                            "5": {},
+                        },
+                    },
+                    "3": {"children": {"6": {}, "7": {}}},
+                },
+            }
+        }
+        actual = export.tree_to_nested_dict_key(binarytree_node)
+        assert actual == expected, f"Expected\n{expected}\nReceived\n{actual}"
+
+
 class TestTreeToDot:
     @staticmethod
     def test_tree_to_dot(binarytree_node):