Test suite for a 'bigtree' library issue

000_bigtree_issue_sample.py

"""
A test suite use for research&development and that demonstrate an issue with
'dict_to_tree'

It has been started with custom Node object sample:

https://bigtree.readthedocs.io/en/latest/others/tips.html#population-node-add-functionality-method-property

Opposed to original sample the percentage attribute have been disabled else it causes
"division by zero" and prevent to make demonstration.

NOTE: Issue has been fixed in bigtree=0.10.3
"""
import json

from bigtree import (
    Node,
    dict_to_tree,
    tree_to_nested_dict,
    print_tree,
    find_child_by_name,
)


EXPECTED_EXPORT = {
    "name": "World",
    "population": 450,
    # "percentage": 1,
    "children": [
        {
            "name": "Country A",
            "population": 150,
            # "percentage": 0.33,
            "children": [
                {
                    "name": "State A1",
                    "population": 100,
                    # "percentage": 0.22
                },
                {
                    "name": "State A2",
                    "population": 50,
                    # "percentage": 0.11
                }
            ]
        },
        {
            "name": "Country B",
            "population": 200,
            # "percentage": 0.44
        },
        {
            "name": "Country C",
            "population": 100,
            # "percentage": 0.22
        }
    ]
}


class PopulationNode(Node):

    def __init__(self, name: str, population: int = 0, **kwargs):
        super().__init__(name, **kwargs)
        self._population = population

    @property
    def population(self):
        if self.is_leaf:
            return self._population
        return sum([child.population for child in self.children])

    @property
    def percentage(self):
        if self.is_root:
            return 1
        return round(self.population / self.root.population, 2)


def test_directly_population_node():
    """
    Tree filled directly with PopulationNode nodes should export into a dictionnary
    as expected.
    """
    root = PopulationNode("World")
    b1 = PopulationNode("Country A", parent=root)
    c1 = PopulationNode("State A1", 100, parent=b1)
    c2 = PopulationNode("State A2", 50, parent=b1)
    b2 = PopulationNode("Country B", 200, parent=root)
    b3 = PopulationNode("Country C", 100, parent=root)

    exported = tree_to_nested_dict(root, attr_dict={
        "population": "population",
    })

    print()
    print("----------------------- show (all_attrs) -----------------------")
    print()
    root.show(all_attrs=True)

    print()
    print("----------------------- show (attr_list) -----------------------")
    print()
    root.show(attr_list=["population"])

    print()
    print("----------------------- print_tree (all_attrs) -----------------------")
    print()
    print_tree(root, all_attrs=True)

    print()
    print("----------------------- print_tree (attr_list) -----------------------")
    print()
    print_tree(root, attr_list=["population"])

    print()
    print("----------------------- dict dump -----------------------")
    print()
    print(json.dumps(exported, indent=4))

    assert exported == EXPECTED_EXPORT


def test_loaded_dict_to_tree():
    """
    Tree filled from a dictionnary with PopulationNode nodes should export into a
    dictionnary as expected.
    """
    path_dict = {
        "World": {"population": 0},
        "World/Country A": {"population": 0},
        "World/Country A/State A1": {"population": 100},
        "World/Country A/State A2": {"population": 50},
        "World/Country B": {"population": 200},
        "World/Country C": {"population": 100},
    }

    root = dict_to_tree(path_dict, node_type=PopulationNode)

    exported = tree_to_nested_dict(root, attr_dict={
        "population": "population",
    })

    print()
    print("----------------------- inspecting nodes -----------------------")
    print()
    print("root:", root)
    print("root.population:", root.population)

    node_country_a = find_child_by_name(root, "Country A")
    print("node_country_a:", node_country_a)

    node_state_a1 = find_child_by_name(node_country_a, "State A1")
    print("node_state_a1:", node_state_a1)
    print("node_state_a1.population:", node_state_a1.population)
    print("node_state_a1.get_attr:", node_state_a1.get_attr("population"))
    print("node_state_a1.describe:", node_state_a1.describe(exclude_prefix="_"))

    print()
    print("----------------------- show (all_attrs) -----------------------")
    print()
    root.show(all_attrs=True)

    print()
    print("----------------------- show (attr_list) -----------------------")
    print()
    root.show(attr_list=["population"])

    print()
    print("----------------------- print_tree (all_attrs) -----------------------")
    print()
    print_tree(root, all_attrs=True)

    print()
    print("----------------------- print_tree (attr_list) -----------------------")
    print()
    print_tree(root, attr_list=["population"])

    print()
    print("----------------------- dict dump -----------------------")
    print()
    print(json.dumps(exported, indent=4))

    # Expected to fail to be able to see debug output
    #assert 1 == 42
    assert exported == EXPECTED_EXPORT

It currently print the following output:

----------------------- inspecting nodes -----------------------

root: PopulationNode(/World, population=0)
root.population: 0
node_country_a: PopulationNode(/World/Country A, population=0)
node_state_a1: PopulationNode(/World/Country A/State A1, population=100)
node_state_a1.population: 0
node_state_a1.get_attr: 0
node_state_a1.describe: [('name', 'State A1'), ('population', 100)]

----------------------- show (all_attrs) -----------------------

World [population=0]
├── Country A [population=0]
│   ├── State A1 [population=100]
│   └── State A2 [population=50]
├── Country B [population=200]
└── Country C [population=100]

----------------------- show (attr_list) -----------------------

World [population=0]
├── Country A [population=0]
│   ├── State A1 [population=0]
│   └── State A2 [population=0]
├── Country B [population=0]
└── Country C [population=0]

----------------------- print_tree (all_attrs) -----------------------

World [population=0]
├── Country A [population=0]
│   ├── State A1 [population=100]
│   └── State A2 [population=50]
├── Country B [population=200]
└── Country C [population=100]

----------------------- print_tree (attr_list) -----------------------

World [population=0]
├── Country A [population=0]
│   ├── State A1 [population=0]
│   └── State A2 [population=0]
├── Country B [population=0]
└── Country C [population=0]

----------------------- dict dump -----------------------

{
    "name": "World",
    "population": 0,
    "children": [
        {
            "name": "Country A",
            "population": 0,
            "children": [
                {
                    "name": "State A1",
                    "population": 0
                },
                {
                    "name": "State A2",
                    "population": 0
                }
            ]
        },
        {
            "name": "Country B",
            "population": 0
        },
        {
            "name": "Country C",
            "population": 0
        }
    ]
}