Reorganise files, create common file (#397)

* 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

Author: kayjan

bigtree/dag/construct.py

@@ -3,7 +3,7 @@
 from typing import Any, Collection, Dict, List, Mapping, Optional, Tuple, Type, TypeVar
 
 from bigtree.node import dagnode
-from bigtree.utils import assertions, exceptions
+from bigtree.utils import assertions, common, exceptions
 
 try:
     import pandas as pd
@@ -201,14 +201,14 @@ def dataframe_to_dag(
     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 = assertions.filter_attributes(
+        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
 
-        if not assertions.isnull(parent_name):
+        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]

bigtree/dag/export.py

@@ -3,7 +3,7 @@
 from typing import Any, Dict, List, Mapping, Optional, Tuple, TypeVar, Union
 
 from bigtree.node import dagnode
-from bigtree.utils import assertions, exceptions, iterators
+from bigtree.utils import assertions, common, exceptions, iterators
 
 try:
     import pandas as pd
@@ -86,28 +86,18 @@ def dag_to_dict(
     for parent_node, child_node in iterators.dag_iterator(dag):
         if parent_node.is_root:
             data_parent: Dict[str, Any] = {}
-            if all_attrs:
-                data_parent.update(
-                    parent_node.describe(
-                        exclude_attributes=["name"], exclude_prefix="_"
-                    )
-                )
-            elif attr_dict:
-                for k, v in attr_dict.items():
-                    data_parent[v] = parent_node.get_attr(k)
+            data_parent = common.assemble_attributes(
+                parent_node, attr_dict, all_attrs, data_parent
+            )
             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]}
-            if all_attrs:
-                data_child.update(
-                    child_node.describe(exclude_attributes=["name"], exclude_prefix="_")
-                )
-            elif attr_dict:
-                for k, v in attr_dict.items():
-                    data_child[v] = child_node.get_attr(k)
+            data_child = common.assemble_attributes(
+                child_node, attr_dict, all_attrs, data_child
+            )
             data_dict[child_node.node_name] = data_child
     return data_dict
 
@@ -155,25 +145,15 @@ 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}
-            if all_attrs:
-                data_parent.update(
-                    parent_node.describe(
-                        exclude_attributes=["name"], exclude_prefix="_"
-                    )
-                )
-            elif attr_dict:
-                for k, v in attr_dict.items():
-                    data_parent[v] = parent_node.get_attr(k)
+            data_parent = common.assemble_attributes(
+                parent_node, attr_dict, all_attrs, data_parent
+            )
             data_list.append(data_parent)
 
         data_child = {name_col: child_node.node_name, parent_col: parent_node.node_name}
-        if all_attrs:
-            data_child.update(
-                child_node.describe(exclude_attributes=["name"], exclude_prefix="_")
-            )
-        elif attr_dict:
-            for k, v in attr_dict.items():
-                data_child[v] = child_node.get_attr(k)
+        data_child = common.assemble_attributes(
+            child_node, attr_dict, all_attrs, data_child
+        )
         data_list.append(data_child)
     return pd.DataFrame(data_list).drop_duplicates().reset_index(drop=True)
 

bigtree/tree/construct/dataframes.py

@@ -5,7 +5,7 @@
 from bigtree.node import node
 from bigtree.tree.construct.dictionaries import add_dict_to_tree_by_name
 from bigtree.tree.construct.strings import add_path_to_tree
-from bigtree.utils import assertions
+from bigtree.utils import assertions, common
 
 try:
     import pandas as pd
@@ -118,7 +118,7 @@ def add_dataframe_to_tree_by_path(
 
     root_node = tree.root
     for row in data.to_dict(orient="index").values():
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             row, omit_keys=["name", path_col], omit_null_values=True
         )
         add_path_to_tree(
@@ -191,7 +191,7 @@ def add_dataframe_to_tree_by_name(
         .to_dict(orient="index")
     )
     name_attrs = {
-        k1: {k2: v2 for k2, v2 in v1.items() if not assertions.isnull(v2)}
+        k1: {k2: v2 for k2, v2 in v1.items() if not common.isnull(v2)}
         for k1, v1 in name_attrs.items()
     }
 
@@ -283,7 +283,7 @@ def add_polars_to_tree_by_path(
 
     root_node = tree.root
     for row_kwargs in data.to_dicts():
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             row_kwargs, omit_keys=["name", path_col], omit_null_values=True
         )
         add_path_to_tree(
@@ -354,7 +354,7 @@ def add_polars_to_tree_by_name(
         .rows_by_key(key=name_col, named=True)
     )
     name_attrs = {
-        k1: {k2: v2 for k2, v2 in v1[0].items() if not assertions.isnull(v2)}
+        k1: {k2: v2 for k2, v2 in v1[0].items() if not common.isnull(v2)}
         for k1, v1 in name_attrs.items()
     }
 
@@ -446,15 +446,15 @@ def dataframe_to_tree(
         root_node_kwargs = list(
             root_node_data[attribute_cols].to_dict(orient="index").values()
         )[0]
-        root_node_kwargs = assertions.filter_attributes(
+        root_node_kwargs = common.filter_attributes(
             root_node_kwargs, omit_keys=["name", path_col], omit_null_values=True
         )
         root_node = node_type(root_name, **root_node_kwargs)
     else:
         root_node = node_type(root_name)
 
     for row in data.to_dict(orient="index").values():
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             row, omit_keys=["name", path_col], omit_null_values=True
         )
         add_path_to_tree(
@@ -562,7 +562,7 @@ def _retrieve_attr(_row: Dict[str, Any]) -> Dict[str, Any]:
         Returns:
             Attribute dictionary
         """
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             _row, omit_keys=[child_col, parent_col], omit_null_values=True
         )
         node_attrs["name"] = _row[child_col]
@@ -678,15 +678,15 @@ def polars_to_tree(
     if len(root_node_data):
         root_node_kwargs_list = root_node_data[attribute_cols].to_dicts()
         root_node_kwargs = root_node_kwargs_list[0] if root_node_kwargs_list else {}
-        root_node_kwargs = assertions.filter_attributes(
+        root_node_kwargs = common.filter_attributes(
             root_node_kwargs, omit_keys=["name", path_col], omit_null_values=True
         )
         root_node = node_type(root_name, **root_node_kwargs)
     else:
         root_node = node_type(root_name)
 
     for row in data.to_dicts():
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             row, omit_keys=["name", path_col], omit_null_values=True
         )
         add_path_to_tree(
@@ -794,7 +794,7 @@ def _retrieve_attr(_row: Dict[str, Any]) -> Dict[str, Any]:
         Returns:
             Attribute dictionary
         """
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             _row, omit_keys=[child_col, parent_col], omit_null_values=True
         )
         node_attrs["name"] = _row[child_col]

bigtree/tree/construct/dictionaries.py

@@ -4,7 +4,7 @@
 
 from bigtree.node import node
 from bigtree.tree.construct.strings import add_path_to_tree
-from bigtree.utils import assertions
+from bigtree.utils import assertions, common
 
 __all__ = [
     "add_dict_to_tree_by_path",
@@ -127,7 +127,7 @@ def add_dict_to_tree_by_name(tree: T, name_attrs: Mapping[str, Mapping[str, Any]
     attr_dict_names = set(name_attrs.keys())
 
     for _node in findall(tree, lambda _node1: _node1.node_name in attr_dict_names):
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             name_attrs[_node.node_name], omit_keys=["name"], omit_null_values=False
         )
         _node.set_attrs(node_attrs)
@@ -201,7 +201,7 @@ def dict_to_tree(
         or path_attrs.get(root_name + sep, {})
         or path_attrs.get(sep + root_name + sep, {})
     )
-    root_node_attrs = assertions.filter_attributes(
+    root_node_attrs = common.filter_attributes(
         root_node_attrs, omit_keys=["name"], omit_null_values=False
     )
     root_node = node_type(
@@ -212,7 +212,7 @@ def dict_to_tree(
 
     # Convert dictionary to dataframe
     for node_path, node_attrs in path_attrs.items():
-        node_attrs = assertions.filter_attributes(
+        node_attrs = common.filter_attributes(
             node_attrs, omit_keys=["name"], omit_null_values=False
         )
         add_path_to_tree(

bigtree/tree/export/dataframes.py

@@ -3,7 +3,7 @@
 from typing import Any, Dict, Optional, TypeVar
 
 from bigtree.node import node
-from bigtree.utils import exceptions
+from bigtree.utils import common, exceptions
 
 try:
     import pandas as pd
@@ -106,13 +106,9 @@ def _recursive_append(_node: T) -> None:
                         parent_name = _node.parent.node_name
                     data_child[parent_col] = parent_name
 
-                if all_attrs:
-                    data_child.update(
-                        _node.describe(exclude_attributes=["name"], exclude_prefix="_")
-                    )
-                elif attr_dict:
-                    for k, v in attr_dict.items():
-                        data_child[v] = _node.get_attr(k)
+                data_child = common.assemble_attributes(
+                    _node, attr_dict, all_attrs, data_child
+                )
                 data_list.append(data_child)
             for _child in _node.children:
                 _recursive_append(_child)
@@ -211,13 +207,9 @@ def _recursive_append(_node: T) -> None:
                         parent_name = _node.parent.node_name
                     data_child[parent_col] = parent_name
 
-                if all_attrs:
-                    data_child.update(
-                        _node.describe(exclude_attributes=["name"], exclude_prefix="_")
-                    )
-                elif attr_dict:
-                    for k, v in attr_dict.items():
-                        data_child[v] = _node.get_attr(k)
+                data_child = common.assemble_attributes(
+                    _node, attr_dict, all_attrs, data_child
+                )
                 data_list.append(data_child)
             for _child in _node.children:
                 _recursive_append(_child)

bigtree/tree/export/dictionaries.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Optional, TypeVar
 
 from bigtree.node import node
+from bigtree.utils import common
 
 __all__ = [
     "tree_to_dict",
@@ -107,7 +108,7 @@ def _recursive_append(_node: T) -> None:
                     if _node.parent:
                         parent_name = _node.parent.node_name
                     data_child[parent_key] = parent_name
-                data_child = _assemble_attributes(
+                data_child = common.assemble_attributes(
                     _node, attr_dict, all_attrs, data_child
                 )
                 data_dict[_node.path_name] = data_child
@@ -165,7 +166,7 @@ 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 = _assemble_attributes(
+                data_child = common.assemble_attributes(
                     _node, attr_dict, all_attrs, data_child
                 )
                 if child_key in parent_dict:
@@ -224,7 +225,7 @@ def _recursive_append(_node: T, parent_dict: Dict[str, Any]) -> None:
         """
         if _node:
             if not max_depth or _node.depth <= max_depth:
-                data_child = _assemble_attributes(_node, attr_dict, all_attrs)
+                data_child = common.assemble_attributes(_node, attr_dict, all_attrs)
                 if child_key in parent_dict:
                     parent_dict[child_key][_node.node_name] = data_child
                 else:

bigtree/tree/export/stdout.py

@@ -3,7 +3,7 @@
 from typing import Any, Collection, Iterable, List, Optional, Tuple, TypeVar, Union
 
 from bigtree.node import node
-from bigtree.utils import assertions, constants
+from bigtree.utils import common, constants
 
 __all__ = [
     "print_tree",
@@ -201,7 +201,7 @@ def print_tree(
                     attr_str_list = [
                         f"{attr_name}={_node.get_attr(attr_name)}"
                         for attr_name in attr_list
-                        if not assertions.isnull(_node.get_attr(attr_name))
+                        if not common.isnull(_node.get_attr(attr_name))
                     ]
                 else:
                     attr_str_list = [

bigtree/utils/assertions.py

@@ -4,7 +4,6 @@
     TYPE_CHECKING,
     Any,
     Collection,
-    Dict,
     Iterable,
     List,
     Mapping,
@@ -44,8 +43,6 @@
     "assert_dataframe_no_duplicate_attribute",
     "assert_dataframe_no_duplicate_children",
     "assert_tree_type",
-    "isnull",
-    "filter_attributes",
 ]
 
 
@@ -252,43 +249,3 @@ def assert_tree_type(
         raise TypeError(
             f"Tree should be of type `{tree_type_name}`, or inherit from `{tree_type_name}`"
         )
-
-
-def isnull(value: Any) -> bool:
-    """Check if value is null.
-
-    Args:
-        value: value to check
-
-    Returns:
-        Flag if value is null
-    """
-    import math
-
-    if value is None or (isinstance(value, float) and math.isnan(value)):
-        return True
-    return False
-
-
-def filter_attributes(
-    node_attributes: Mapping[str, Any],
-    omit_keys: Collection[str],
-    omit_null_values: bool,
-) -> Dict[str, Any]:
-    """Filter node attributes to remove certain keys and/or values.
-
-    Args:
-        node_attributes: node attributes information
-        omit_keys: keys to omit
-        omit_null_values: indicator whether to omit values that are null
-
-    Returns:
-        Filtered node attributes
-    """
-    if omit_null_values:
-        return {
-            k: v
-            for k, v in node_attributes.items()
-            if not isnull(v) and k not in omit_keys
-        }
-    return {k: v for k, v in node_attributes.items() if k not in omit_keys}