construct_tree.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   refresh_tree.py
@Time    :   2022/01/20 18:10:12
@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

from .utils import read_json
from .utils import write_json
from .tree_common import TreeCommon


class ConstructTree:
    def __init__(self):
        self.tree = Tree()

    def load(self):
        pass

    def construct_tree_from_json(self):
        json_file_name = "./data/tree.json"
        tree_dict = read_json(json_file_name)
        TreeCommon.construct_tree(self.tree, tree_dict, None)
    
    def _create_config_file(self, node : Node, dir, node_level):
        file_path = os.path.join(dir, "config.json")
        if node_level < 3:
            data = {
                "keywords": node.data["keywords"],
                "node_id": node.identifier,
                "keywords_must": node.data["keywords_must"],
                "keywords_forbid": node.data["keywords_forbid"]
            }
        else:
            data = TreeCommon.get_tree_json_recursively(self.tree, node.identifier)
            data = list(data.values())[0]
            data["export"] = []

        write_json(file_path, data)

    def create_contents(self):
        # 先序遍历
        def _preorder_traversal_create_contents(cur_node : Node, dir, index=1):
            ''' 递归：先根遍历所有结点 '''
            
            cur_node_id = cur_node.identifier
            cur_node_tag = cur_node.tag
            cur_node_level = self.tree.level(cur_node_id)
            all_children = self.tree.children(cur_node_id)
            if cur_node_level < 4:
                if not cur_node.is_root():
                    dir = os.path.join(dir, str(index) + "." + cur_node_tag)
                    if not os.path.exists(dir):
                        os.mkdir(dir)
                self._create_config_file(cur_node, dir, cur_node_level)
                for index, child_node in enumerate(all_children):
                    _preorder_traversal_create_contents(child_node, dir, index + 1)

        root_node_id = self.tree.root
        root_node = self.tree.get_node(root_node_id)
        root_dir = "./data"
        _preorder_traversal_create_contents(root_node, root_dir)