tree_common.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   tree_common.py
@Time    :   2022/01/20 18:17:41
@Author  :   Lu Xin 
@Contact :   luxin@csdn.net
'''

# here put the import lib
import os
import ipdb
from tqdm import tqdm
from treelib import Tree
from treelib import Node


class TreeCommon:
    def __init__(self):
        pass

    def load(self):
        pass

    @classmethod
    def construct_tree(cls, tree : Tree, tree_dict, parent):
        # tree_name = list(tree_dict.keys())[0]
        # root_node = Node(
        #     tag="tree_name", 
        #     identifier=tree_dict[tree_name]["node_id"], 
        #     data={
        #         "keywords": tree_dict[tree_name]["keywords"],
        #         "keywords_must": tree_dict[tree_name]["keywords_must"],
        #         "keywords_forbid": tree_dict[tree_name]["keywords_forbid"]
        # })
        # tree.add_node(root_node, parent=None)

        for node_text, node_info in tree_dict.items():
            node_id = node_info["node_id"]
            node_kw = node_info["keywords"]
            node_kw_must = node_info["keywords_must"]
            node_kw_forbid = node_info["keywords_forbid"]
            subtree_list = node_info["children"]
            node = Node(
                tag=node_text, 
                identifier=node_id, 
                data={
                    "keywords": node_kw,
                    "keywords_must": node_kw_must,
                    "keywords_forbid": node_kw_forbid
                })
            tree.add_node(node, parent=parent)
            for subtree_dict in subtree_list:
                cls.construct_tree(tree, subtree_dict, node_id)

    @classmethod
    def get_tree_json_recursively(cls, tree : Tree, cur_node_id):
        struct = {}

        cur_node = tree.get_node(cur_node_id)
        cur_node_text = cur_node.tag

        struct[cur_node_text] = {
            "keywords": cur_node.data["keywords"],
            "keywords_must": cur_node.data["keywords_must"],
            "keywords_forbid": cur_node.data["keywords_forbid"],
            "children": [],
            "node_id": cur_node_id
        }

        all_children_list = [m_node.identifier for m_node in tree.children(cur_node_id)]

        for child_node in all_children_list:
            child_struct = cls.get_tree_json_recursively(tree, child_node)
            struct[cur_node_text]["children"].append(child_struct)

        return struct