dbms_sanya/python/plugins/ResourceTree/ui/TreeWidget.py

from typing import Optional
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem
from .PostgreSQL import PostgreSQL


class Tree:
    type = ""

    def __init__(self, type: Optional[str] = ''):
        self.type = type

    def createTree(self, data, parent, parentid):
        """递归创建树形结构"""
        # 先获取当前层级的所有节点并按sort排序
        current_level_nodes = []
        for row in data:
            if (parentid == '' and row['pbsm'] == '') or (row['pbsm'] == parentid):
                current_level_nodes.append(row)

        # 按sort字段排序
        current_level_nodes.sort(key=lambda x: int(x['sort']) if x['sort'] is not None else 9999)

        # 创建当前层级的节点
        for row in current_level_nodes:
            item = QTreeWidgetItem(parent)
            item.setText(0, row['name'])       # 第0列：名称
            item.setText(1, row['bsm'])        # 第1列：唯一ID

            # 设置第2列：图层类型（如 vector、raster）或目录的 pbsm
            if 'layertype' in row and row['layertype']:
                item.setText(2, row['layertype'])  # 数据节点
            else:
                item.setText(2, row['pbsm'])       # 目录节点

            # 用于排序显示
            item.setData(0, Qt.UserRole + 1, int(row['sort']) if row['sort'] is not None else 9999)

            # 递归创建子节点
            self.createTree(data, item, row['bsm'])

    def initTreeWidget(self, dbcoon):
        """初始化QTreeWidget"""
        try:
            print("初始化 QTreeWidget...")
            tree = QTreeWidget()
            tree.setStyleSheet("QTreeView::item{padding: 2px;}")
            tree.setColumnCount(3)
            tree.setColumnHidden(1, True)
            tree.setColumnHidden(2, True)
            tree.setObjectName(self.type)
            tree.setHeaderLabels(["名称", "实体数据名称", "数据类型"])
            tree.headerItem().setHidden(True)

            print("准备连接数据库...")
            pg = PostgreSQL(schema=dbcoon["schema"], host=dbcoon["host"],
                            port=dbcoon["port"],
                            user=dbcoon["user"],
                            password=dbcoon["password"],
                            dbname=dbcoon["dbname"])
            print("数据库连接成功")

            print("获取资源目录数据...")
            treedata = pg.getZyml()  # 你需要确保这里返回的数据中包含 layertype 字段
            print(f"资源目录数据数量: {len(treedata)}")

            self.createTree(treedata, tree, "")
            print("树结构创建成功")

            pg.close()
            print("数据库连接关闭")

            return tree
        except Exception as e:
            import traceback
            print("initTreeWidget 初始化失败:", e)
            print(traceback.format_exc())
            raise e  # 抛出给外层捕获