Rong & Wei

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sqlite3
from PyQt5.QtWidgets import QTreeView
from PyQt5.QtCore import QAbstractItemModel, QModelIndex, Qt, QVariant
from PyQt5.QtGui import QColor

ACCOUNT_TREE = {
    ('表1', '资产评估清查结果--汇总表'): None,
    ('表2', '资产评估清查结果分类汇总表'): None,
    ('表3', '流动资产评估清查汇总表'): {
        ('表3-1', '货币资金评估汇总表'): {
            ('表3-1-1', '货币资金—现金评估清查明细表'): None,
            ('表3-1-2', '货币资金—银行存款评估清查明细表'): None,
            ('表3-1-3', '货币资金—其他货币资金评估清查明细表'): None},
        ('表3-2', '交易性金融资产评估清查汇总表'): {
            ('表3-2-1', '交易性金融资产—股票投资评估清查明细表'): None,
            ('表3-2-2', '交易性金融资产—债券投资评估清查明细表'): None,
            ('表3-2-3', '交易性金融资产—基金投资评估清查明细表'): None},
        ('表3-3', '应收票据评估清查明细表'): None,
        ('表3-4', '应收账款评估清查明细表'): None,
        ('表3-5', '预付账款评估清查明细表'): None,
        ('表3-6', '应收利息评估清查明细表'): None,
        ('表3-7', '应收股利（应收利润）评估清查明细表'): None,
        ('表3-8', '其他应收款评估清查明细表'): None,
        ('表3-9', '存货评估汇总表'): {
            ('表3-9-1', '存货—材料采购（在途物资）评估清查明细表'): None,
            ('表3-9-2', '存货—原材料评估清查明细表'): None,
            ('表3-9-3', '存货—在库周转材料评估清查明细表'): None,
            ('表3-9-4', '存货—委托外加工物资评估清查明细表'): None,
            ('表3-9-5', '存货—产成品（库存商品、开发产品、农产品）评估清查明细表'): None,
            ('表3-9-6', '存货—在产品（自制半成品）评估清查明细表'): None,
            ('表3-9-7', '存货—发出商品评估清查明细表'): None,
            ('表3-9-8', '存货—在用周转材料评估清查明细表'): None},
        ('表3-10', '一年内到期的非流动资产评估清查明细表'): None,
        ('表3-11', '其他流动资产评估清查明细表'): None},
    ('表4', '非流动资产评估清查汇总表'): {
        ('表4-1', '可供出售金融资产评估清查汇总表'): {
            ('表4-1-1', '可供出售金融资产-股票投资评估清查明细表'): None,
            ('表4-1-2', '可供出售金融资产-债券投资评估清查明细表'): None,
            ('表4-1-3', '可供出售金融资产-其他投资评估清查明细表'): None},
        ('表4-2', '持有至到期投资评估清查明细表'): None,
        ('表4-3', '长期应收款评估清查明细表'): None,
        ('表4-4', '长期股权投资评估清查明细表'): None,
        ('表4-5', '投资性房地产评估汇总表'): {
            ('表4-5-1', '投资性房地产-房屋评估清查明细表'): None,
            ('表4-5-2', '投资性房地产-土地使用权评估清查明细表'): None},
        ('表4-6', '固定资产评估清查汇总表'): {
            ('表4-6-1', '固定资产—房屋建筑物评估清查明细表'): None,
            ('表4-6-2', '固定资产—构筑物及其他辅助设施评估清查明细表'): None,
            ('表4-6-3', '固定资产—管道和沟槽评估清查明细表'): None,
            ('表4-6-4', '固定资产—机器设备评估清查明细表'): None,
            ('表4-6-5', '固定资产—车辆评估清查明细表'): None,
            ('表4-6-6', '固定资产—电子设备评估清查明细表'): None,
            ('表4-6-7', '固定资产—土地评估清查明细表'): None},
        ('表4-7', '在建工程评估汇总表'): {
            ('表4-7-1', '在建工程—土建工程评估清查明细表'): None,
            ('表4-7-2', '在建工程—设备安装工程评估清查明细表'): None},
        ('表4-8', '工程物资评估清查明细表'): None,
        ('表4-9', '固定资产清理评估清查明细表'): None,
        ('表4-10', '生产性生物资产评估清查明细表'): None,
        ('表4-11', '油气资产评估清查明细表'): None,
        ('表4-12', '无形资产评估清查汇总表'): {
            ('表4-12-1', '无形资产—土地使用权评估清查明细表'): None,
            ('表4-12-2', '无形资产-矿业权评估清查明细表'): None,
            ('表4-12-3', '无形资产—其他无形资产评估清查明细表'): None},
        ('表4-13', '开发支出评估清查明细表'): None,
        ('表4-14', '商誉评估清查明细表'): None,
        ('表4-15', '长期待摊费用评估清查明细表'): None,
        ('表4-16', '递延所得税资产评估清查明细表'): None,
        ('表4-17', '其他非流动资产评估清查明细表'): None},
    ('表5', '流动负债评估汇总表'): {
        ('表5-1', '短期借款评估清查明细表'): None,
        ('表5-2', '交易性金融负债评估清查明细表'): None,
        ('表5-3', '应付票据评估清查明细表'): None,
        ('表5-4', '应付账款评估清查明细表'): None,
        ('表5-5', '预收账款评估清查明细表'): None,
        ('表5-6', '应付职工薪酬评估清查明细表'): None,
        ('表5-7', '应交税费评估清查明细表'): None,
        ('表5-8', '应付利息评估清查明细表'): None,
        ('表5-9', '应付股利（应付利润）评估清查明细表'): None,
        ('表5-10', '其他应付款评估清查明细表'): None,
        ('表5-11', '一年内到期的非流动负债评估清查明细表'): None,
        ('表5-12', '其他流动负债评估清查明细表'): None},
    ('表6', '非流动负债评估汇总表'): {
        ('表6-1', '长期借款评估清查明细表'): None,
        ('表6-2', '应付债券评估清查明细表'): None,
        ('表6-3', '长期应付款评估清查明细表'): None,
        ('表6-4', '专项应付款评估清查明细表'): None,
        ('表6-5', '预计负债评估清查明细表'): None,
        ('表6-6', '递延所得税负债评估清查明细表'): None,
        ('表6-7', '其他非流动负债评估清查明细表'): None},
    ('表7', '所有者权益表'): None,
    ('表8', '资产减值准备备查付表'): None}


class TreeItem:
    def __init__(self, row_data, parent=None):
        self.parent_item = parent
        self.item_data = row_data
        self.child_items = []

    def append_child(self, item):
        self.child_items.append(item)

    def child(self, row):
        return self.child_items[row]

    def child_count(self):
        return len(self.child_items)

    def column_count(self):
        return len(self.item_data)

    def data(self, column):
        try:
            return self.item_data[column]
        except IndexError:
            return QVariant()

    def parent(self):
        return self.parent_item

    def row(self):
        if self.parent_item:
            return self.parent_item.child_items.index(self)
        return 0


class AccountModel(QAbstractItemModel):
    def __init__(self, conn: sqlite3.Connection, account_tree: dict, parent=None):
        super(AccountModel, self).__init__(parent)
        self.conn = conn
        self.c = self.conn.cursor()
        self.root_item = TreeItem(("编号", "名称"))
        self.add_data(account_tree, self.root_item)

    def columnCount(self, parent=None, *args, **kwargs):
        if parent.isValid():
            return parent.internalPointer().column_count()
        else:
            return self.root_item.column_count()

    def data(self, index: QModelIndex, role=None):
        if not index.isValid():
            return QVariant()
        if role == Qt.TextColorRole:
            table_name = index.internalPointer().data(0)
            try:
                if self.c.execute(f'SELECT count(*) FROM [{table_name}];').fetchone()[0]:
                    return QColor(Qt.blue)
            except sqlite3.OperationalError as e:
                print(e)
        if role != Qt.DisplayRole:
            return QVariant()
        item = index.internalPointer()
        return item.data(index.column())

    @staticmethod
    def code(index: QModelIndex, role=None):
        if not index.isValid():
            return QVariant()
        if role != Qt.DisplayRole:
            return QVariant()
        item = index.internalPointer()
        return item.data(0)

    def flags(self, index: QModelIndex):
        if not index.isValid():
            return Qt.NoItemFlags
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def headerData(self, section, orientation, role=None):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.root_item.data(section)
        return QVariant()

    def index(self, row, column, parent=None, *args, **kwargs):
        if not self.hasIndex(row, column, parent):
            return QModelIndex()
        if not parent.isValid():
            parent_item = self.root_item
        else:
            parent_item = parent.internalPointer()
        child_item = parent_item.child(row)
        if child_item:
            return self.createIndex(row, column, child_item)
        else:
            return QModelIndex()

    def parent(self, index: QModelIndex=None):
        if not index.isValid():
            return QModelIndex()
        child_item = index.internalPointer()
        parent_item = child_item.parent()
        if parent_item == self.root_item:
            return QModelIndex()
        return self.createIndex(parent_item.row(), 0, parent_item)

    def rowCount(self, parent=None, *args, **kwargs):
        if parent.column() > 0:
            return 0
        if not parent.isValid():
            parent_item = self.root_item
        else:
            parent_item = parent.internalPointer()
        return parent_item.child_count()

    def add_data(self, data, parent):
        for rid, ((code, name), setting) in enumerate(data.items()):
            parent.append_child(TreeItem([code, name], parent))
            if setting:
                self.add_data(setting, parent.child(rid))


class AccountTree(QTreeView):
    def __init__(self, conn, parent=None):
        super(AccountTree, self).__init__(parent)
        self.header().setDefaultAlignment(Qt.AlignCenter)
        self.setExpandsOnDoubleClick(False)
        self.horizontalScrollBar().setStyleSheet('QScrollBar:horizontal{height:0px;}')
        account_model = AccountModel(conn, ACCOUNT_TREE)
        self.setModel(account_model)