Skip to content

Instantly share code, notes, and snippets.

@sveetch
Last active August 12, 2023 11:12
Show Gist options
  • Save sveetch/7d60f0cfcc76ae31f2ca973038da3c88 to your computer and use it in GitHub Desktop.
Save sveetch/7d60f0cfcc76ae31f2ca973038da3c88 to your computer and use it in GitHub Desktop.
Test suite for a 'bigtree' library issue
"""
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
@sveetch
Copy link
Author

sveetch commented Aug 11, 2023

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
        }
    ]
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment