Wrapper over Treelib to easily generate and update required data for nicegui Tree

TreeView.py

from typing import Dict
from treelib import Node, Tree

class TreeNode:
    """If parent_id is -1, node at root level"""
    def __init__(self, id, parent_id, disp_text, node_data = None) -> None:
        self.id: int = id
        self.parent_id: int = parent_id
        self.disp_text: str = disp_text
        self.node_data = node_data

    def to_dict(self):
        return {'id': self.id, 'parent_id': self.parent_id, 'disp_text': self.disp_text, 'node_data': self.node_data}

class TreeView:
    ID_FIELD = 'id'
    LABEL_FIELD = 'label'
    CHILDREN_FIELD = 'children'

    def __init__(self, root_label: str = None) -> None:
        # self.on_node_click_cb = on_node_click_cb

        self.root_label: str = root_label
        self.root_nid: str = None

        # Map node id to node, so that node look-up is fast
        self.tree_data: Dict[str, TreeNode] = {}
        pass

    def add_node(self, id: int, disp_text: str, parent_id: int = -1, node_data = None):
        self.tree_data[id] = TreeNode(id=id, disp_text=disp_text, parent_id=parent_id, node_data=node_data)

    def get_node(self, id):
        if isinstance(id, int):
            return self.tree_data[id]
        
        # In other cases, return None
        return None
            

    def _get_node_content(self, node: TreeNode):
        return {self.ID_FIELD: node.id, self.LABEL_FIELD: node.disp_text, self.CHILDREN_FIELD: []}

    def _build_tree(self):
        tree = Tree()

        root = 'tree_root' if self.root_label is None else self.root_label
        self.root_nid = root

        #Add virtual parent which is to ensure exactly 1 root
        tree.create_node(root, root, parent=None, data=None)

        # First pass, create all nodes under root
        for node in self.tree_data.values():
            if node.id in tree:
                print('Cant add', node.id)
            else:
                tree.create_node(tag=node.disp_text, identifier=node.id, parent=root, data=node)

        # tree.show()

        # Now update parent
        for node in tree.all_nodes_itr():
            if node.identifier == root:
                continue
            nd: TreeNode = node.data
            parent = root if nd.parent_id == -1 else nd.parent_id
            tree.move_node(node.identifier, parent)

        # tree.show()
        return tree

    def _get_tree_node(self, nid: str):
        # TODO: get children for nid

        node: Node = self.tree_content.get_node(nid)

        if node is None:
            return None
        
        if nid == self.root_nid:
            node_dict = {self.ID_FIELD: nid, self.LABEL_FIELD: self.root_label, self.CHILDREN_FIELD: []}
        else:
            node_dict = self._get_node_content(node.data)

        children = self.tree_content.children(nid)

        for ch_node in children:
            ch_id = ch_node.identifier
            ch_children = children = self.tree_content.children(ch_id)
            if len(ch_children) > 0:
                # Recurse...
                ch_data = self._get_tree_node(ch_id)
                if ch_data is None:
                    raise ValueError('Something went wrong')
                
                node_dict['children'].append(ch_data)
            else:
                node_dict['children'].append(self._get_node_content(ch_node.data))
        
        return node_dict

    def get_nicegui_treedata(self):
        # Build the tree as required by nicegui
        self.tree_content = self._build_tree()

        return [self._get_tree_node(self.root_nid)]

# Test
en_test = False

if en_test:
    from nicegui import ui

    tree_view = TreeView(root_label='My Digital Vault')

    tree_view.add_node(id=0, parent_id=-1, disp_text='L1a')
    tree_view.add_node(id=1, parent_id=0, disp_text='L2a')
    tree_view.add_node(id=2, parent_id=1, disp_text='L3a')
    tree_view.add_node(id=3, parent_id=0, disp_text='L2b')

    tree_view.add_node(id=4, parent_id=-1, disp_text='L1b')

    ui.label("Tree")

    data = tree_view.get_nicegui_treedata()
    ui.tree(data, label_key='label', on_select=lambda e: ui.notify(tree_view.get_node(e.value)))

    ui.run()

Hey, I wanted to ask under which license you publish the gist. Is it okay for you if I use parts of the code in my project for a thesis at university for research purposes?