diff --git a/tests/test_tree.py b/tests/test_tree.py index 43f5fc1..9894d3e 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -712,9 +712,12 @@ def test_shallow_paste(self): n3 = t3.create_node(identifier="C") t1.paste(n1.identifier, t2) - self.assertEqual(t1.to_dict(), {"A": {"children": ["B"]}}) + self.assertEqual(t1.to_dict(), {"A": {"id": "A", "children": [{"B": {"id": "B"}}]}}) t1.paste(n1.identifier, t3) - self.assertEqual(t1.to_dict(), {"A": {"children": ["B", "C"]}}) + self.assertEqual( + t1.to_dict(), + {"A": {"id": "A", "children": [{"B": {"id": "B"}}, {"C": {"id": "C"}}]}}, + ) self.assertEqual(t1.level(n1.identifier), 0) self.assertEqual(t1.level(n2.identifier), 1) diff --git a/treelib/tree.py b/treelib/tree.py index c2594f3..2213b09 100644 --- a/treelib/tree.py +++ b/treelib/tree.py @@ -237,6 +237,26 @@ def __init__(self, tag, identifier=None): if tree.identifier != self._identifier: new_node.clone_pointers(tree.identifier, self._identifier) + @classmethod + def from_json(cls, raw: Union[str, bytes, bytearray]): + """ + Load tree from exported JSON. + """ + tree = cls() + json_parsed = json.loads(raw) + + def _append_node(subtree, parent_id=None): + for tag, node_info in subtree.items(): + node_id = node_info["id"] + node_data = node_info.get("data") + tree.create_node(tag=tag, identifier=node_id, parent=parent_id, data=node_data) + + for child in node_info.get("children", []): + _append_node(child, parent_id=node_id) + + _append_node(json_parsed) + return tree + def _clone( self, identifier: Optional[str] = None, @@ -1960,7 +1980,8 @@ def to_dict( nid = self.root if (nid is None) else nid ntag = self[nid].tag - tree_dict = {ntag: {"children": []}} + tree_dict = {ntag: {"id": nid, "children": []}} + if with_data: tree_dict[ntag]["data"] = self[nid].data @@ -1975,7 +1996,8 @@ def to_dict( self.to_dict(elem.identifier, with_data=with_data, sort=sort, reverse=reverse) ) if len(tree_dict[ntag]["children"]) == 0: - tree_dict = self[nid].tag if not with_data else {ntag: {"data": self[nid].data}} + tree_dict = {ntag: {"id": nid}} if not with_data else {ntag: {"id": nid, "data": self[nid].data}} + return tree_dict def to_json(self, with_data: bool = False, sort: bool = True, reverse: bool = False):