feat: more classes

change pg and category to class

src/budget/__init__.py

@@ -1,15 +1,13 @@
 from PyQt6.QtWidgets import *
 from PyQt6.QtCore import *
 
-class MainWindow(QMainWindow):
+from ._pg import PostgresTab
 
-    cur = None
+class MainWindow(QMainWindow):
-    conn = None
 
     def __init__(self):
         super().__init__()
         self.initUI()
     
-    from ._mw import initUI
+    from ._mw import initUI, onTabChanged
-    from ._pg import connectToDatabase, disconnectFromDatabase, onConnectClicked, initConnectTab
+    from ._category import CategoryTab
-    from ._category import initCategorytTab, onCategoryAddClicked

src/budget/_category.py

@@ -1,30 +1,33 @@
 from PyQt6.QtWidgets import *
+from ._tab import TabPage
+from ._pg import PostgresTab
 import json
 
-def initCategorytTab(self):
+class CategoryTab(TabPage):
-    # 创建连接信息控件
+    def __init__(self, parent):
-    self.categoryComboBox = QComboBox()
+        super().__init__()
+        self.initUI()
+        self.pg = parent.pg
+        self.rows = []
 
-    self.categoryAddButton = QPushButton('Add New Category')
+    def initUI(self):
-    self.categoryAddButton.clicked.connect(self.onCategoryAddClicked)
+        # 创建连接信息控件
+        self.categoryComboBox = QComboBox()
         
-    # 创建控件布局
+        self.categoryAddButton = QPushButton('Add New Category')
-    categoryLayout = QHBoxLayout()
+        self.categoryAddButton.clicked.connect(self.onCategoryAddClicked)
-    categoryLayout.addWidget(self.categoryComboBox)
-    categoryLayout.addWidget(self.categoryAddButton)
 
-    # 连接信号槽
+        # 创建控件布局
-    # self.connectButton.clicked.connect(self.onConnectClicked)
+        categoryLayout = QHBoxLayout()
+        categoryLayout.addWidget(self.categoryComboBox)
+        categoryLayout.addWidget(self.categoryAddButton)
 
-    categoryWidget = QWidget()
+        categoryWidget = QWidget()
-    categoryWidget.setLayout(categoryLayout)
+        self.setLayout(categoryLayout)
-    return categoryWidget
 
 
-def onCategoryAddClicked(self):
+    def onCategoryAddClicked(self):
-    try:
+
-        if (self.cur is None):
-            raise Exception('Not connected to database')
 
         # 创建添加数据对话框
         addDialog = QDialog(self)
@@ -76,14 +79,29 @@ def onCategoryAddClicked(self):
                 data = {"name": line1.text(), "income_type": type, "description": description}
                 data = json.dumps(data)
                 print("insert data:",data)
-                self.cur.execute("INSERT INTO category (meta) VALUES (%s)", (data,))
+                
-                self.conn.commit()
+                self.pg.execute("INSERT INTO category (meta) VALUES (%s)", (data,))
+
+                # 刷新表格
+                self.selected()
 
             except Exception as e:
                 # 处理插入数据错误
                 print(e)
                 QMessageBox.critical(self, 'Error', str(e))
-    except Exception as e:
+
-        # 处理获取表格列名错误
+
-        print(e)
+    def selected(self):
-        QMessageBox.critical(self, 'Error', str(e))
+        # 清空表格
+        self.categoryComboBox.clear()
+        # 获取表格列名
+        try:
+            self.pg.execute("SELECT c_id, meta ->> 'name' FROM category ORDER BY c_id")
+            self.rows = self.pg.cur.fetchall()
+            print(self.rows)
+            for row in self.rows:
+                data = row[1]
+                self.categoryComboBox.addItem(data)
+        except Exception as e:
+            print(e)
+            QMessageBox.critical(self, 'Error', str(e))

src/budget/_mw.py

@@ -1,18 +1,27 @@
 from PyQt6.QtWidgets import *
-from ._pg import initConnectTab
+from ._pg import PostgresTab
-from ._category import initCategorytTab
+from ._category import CategoryTab
+from ._pg import PostgresTab
 
 def initUI(self):
+
     self.setWindowTitle('budget')
     self.setGeometry(100, 100, 800, 600)
 
-    connectTab = initConnectTab(self)
+    self.pg = PostgresTab()
-    categoryTab = initCategorytTab(self)
+    self.categoryTab = CategoryTab(self)
 
     # 创建主窗口tab布局
     mainWidget = QTabWidget()
-    mainWidget.addTab(connectTab, 'Connect')
+    mainWidget.addTab(self.pg, 'Connect')
-    mainWidget.addTab(categoryTab, 'Category')
+    mainWidget.addTab(self.categoryTab, 'Category')
+
+    mainWidget.currentChanged.connect(self.onTabChanged)
 
     self.setCentralWidget(mainWidget)
 
+    self.pg.onConnectClicked()
+
+def onTabChanged(self, index):
+    if (index == 1):
+        self.centralWidget().widget(index).selected()

src/budget/_pg.py

@@ -1,84 +1,114 @@
 from PyQt6.QtWidgets import *
 import psycopg2
+from ._tab import TabPage
 
-def connectToDatabase(self):
+class PostgresTab(TabPage):
-    try:
+    def __init__(self):
-        self.conn = psycopg2.connect(database=self.dbName, user=self.user, password=self.password, host=self.host, port=self.port)
+        super().__init__()
-        self.cur = self.conn.cursor()
+        self.initUI()
-    except Exception as e:
-        print(e)
-        self.cur = None
-        QMessageBox.critical(self, 'Error', str(e))
 
-def disconnectFromDatabase(self):
+    cur = None
-    try:
+    conn = None
-        if (self.cur is None):
+
-            raise Exception('Not connected to database')
+    def connected(self):
-        self.cur.close()
+        return self.cur is not None
-        self.conn.close()
-        self.cur = None
-    except Exception as e:
-        print(e)
-        QMessageBox.critical(self, 'Error', str(e))
         
 
-def onConnectClicked(self):
+    def execute(self,text,values=None):
-    if self.connectButton.text() == 'Connect':
+        try:
-        self.host = self.hostLineEdit.text()
+            if (self.cur is None):
-        self.port = self.portLineEdit.text()
+                raise Exception('Not connected to database')
-        self.dbName = self.dbNameLineEdit.text()
+            self.cur.execute(text,values)
-        self.user = self.userLineEdit.text()
+            self.conn.commit()
-        self.password = self.passwordLineEdit.text()
+        except Exception as e:
-        # 连接到 PostgreSQL 数据库
+            print(e)
-        self.connectToDatabase()
+            QMessageBox.critical(self, 'Error', str(e))
-        self.connectButton.setText('Disconnect')
+
-    else:
+    def fetchall(self):
-        self.disconnectFromDatabase()
+        try:
-        self.connectButton.setText('Connect')
+            if (self.cur is None):
+                raise Exception('Not connected to database')
+            return self.cur.fetchall()
+        except Exception as e:
+            print(e)
+            QMessageBox.critical(self, 'Error', str(e))
+
+    def connectToDatabase(self):
+        try:
+            self.conn = psycopg2.connect(database=self.dbName, user=self.user, password=self.password, host=self.host, port=self.port)
+            self.cur = self.conn.cursor()
+        except Exception as e:
+            print(e)
+            self.cur = None
+            QMessageBox.critical(self, 'Error', str(e))
+
+    def disconnectFromDatabase(self):
+        try:
+            if (self.cur is None):
+                raise Exception('Not connected to database')
+            self.cur.close()
+            self.conn.close()
+            self.cur = None
+        except Exception as e:
+            print(e)
+            QMessageBox.critical(self, 'Error', str(e))
 
 
-def initConnectTab(self):
+    def onConnectClicked(self):
-    # 创建连接信息控件
+        if self.connectButton.text() == 'Connect':
-    self.hostLabel = QLabel('Host:')
+            self.host = self.hostLineEdit.text()
-    self.hostLineEdit = QLineEdit()
+            self.port = self.portLineEdit.text()
-    self.hostLineEdit.setText('159.75.75.169')
+            self.dbName = self.dbNameLineEdit.text()
+            self.user = self.userLineEdit.text()
+            self.password = self.passwordLineEdit.text()
+            # 连接到 PostgreSQL 数据库
+            self.connectToDatabase()
+            self.connectButton.setText('Disconnect')
+        else:
+            self.disconnectFromDatabase()
+            self.connectButton.setText('Connect')
 
-    self.portLabel = QLabel('Port:')
-    self.portLineEdit = QLineEdit()
-    self.portLineEdit.setText('5432')
 
-    self.dbNameLabel = QLabel('Database:')
+    def initUI(self):
-    self.dbNameLineEdit = QLineEdit()
+        # 创建连接信息控件
-    self.dbNameLineEdit.setText('budget')
+        self.hostLabel = QLabel('Host:')
+        self.hostLineEdit = QLineEdit()
+        self.hostLineEdit.setText('159.75.75.169')
 
-    self.userLabel = QLabel('Username:')
+        self.portLabel = QLabel('Port:')
-    self.userLineEdit = QLineEdit()
+        self.portLineEdit = QLineEdit()
-    self.userLineEdit.setText('budget')
+        self.portLineEdit.setText('5432')
 
-    self.passwordLabel = QLabel('Password:')
+        self.dbNameLabel = QLabel('Database:')
-    self.passwordLineEdit = QLineEdit()
+        self.dbNameLineEdit = QLineEdit()
-    self.passwordLineEdit.setEchoMode(QLineEdit.EchoMode.Password)
+        self.dbNameLineEdit.setText('budget')
-    self.passwordLineEdit.setText('budget')
 
-    self.connectButton = QPushButton('Connect')
+        self.userLabel = QLabel('Username:')
+        self.userLineEdit = QLineEdit()
+        self.userLineEdit.setText('budget')
 
-    # 连接信号槽
+        self.passwordLabel = QLabel('Password:')
-    self.connectButton.clicked.connect(self.onConnectClicked)
+        self.passwordLineEdit = QLineEdit()
+        self.passwordLineEdit.setEchoMode(QLineEdit.EchoMode.Password)
+        self.passwordLineEdit.setText('budget')
 
-    # 创建连接信息控件布局
+        self.connectButton = QPushButton('Connect')
-    connectLayout = QHBoxLayout()
-    connectLayout.addWidget(self.hostLabel)
-    connectLayout.addWidget(self.hostLineEdit)
-    connectLayout.addWidget(self.portLabel)
-    connectLayout.addWidget(self.portLineEdit)
-    connectLayout.addWidget(self.dbNameLabel)
-    connectLayout.addWidget(self.dbNameLineEdit)
-    connectLayout.addWidget(self.userLabel)
-    connectLayout.addWidget(self.userLineEdit)
-    connectLayout.addWidget(self.passwordLabel)
-    connectLayout.addWidget(self.passwordLineEdit)
-    connectLayout.addWidget(self.connectButton)
 
-    connectWidget = QWidget()
+        # 连接信号槽
-    connectWidget.setLayout(connectLayout)
+        self.connectButton.clicked.connect(self.onConnectClicked)
-    return connectWidget
+
+        # 创建连接信息控件布局
+        connectLayout = QHBoxLayout()
+        connectLayout.addWidget(self.hostLabel)
+        connectLayout.addWidget(self.hostLineEdit)
+        connectLayout.addWidget(self.portLabel)
+        connectLayout.addWidget(self.portLineEdit)
+        connectLayout.addWidget(self.dbNameLabel)
+        connectLayout.addWidget(self.dbNameLineEdit)
+        connectLayout.addWidget(self.userLabel)
+        connectLayout.addWidget(self.userLineEdit)
+        connectLayout.addWidget(self.passwordLabel)
+        connectLayout.addWidget(self.passwordLineEdit)
+        connectLayout.addWidget(self.connectButton)
+
+        self.setLayout(connectLayout)

src/budget/_tab.py

@@ -0,0 +1,10 @@
+from PyQt6.QtWidgets import *
+import psycopg2
+
+class TabPage(QWidget):
+
+    def __init__(self):
+        super().__init__()
+
+    def selected(self):
+        pass