Restructure UI code and clean up PandasModel

banking_breakdown/ui/__init__.py

@@ -0,0 +1,17 @@
+import sys
+
+import pandas as pd
+from PyQt6.QtWidgets import QApplication
+
+from banking_breakdown.ui.main_window import MainWindow
+
+
+def show_main_window(df: pd.DataFrame = None):
+    app = QApplication(sys.argv)
+    window = MainWindow()
+
+    if df is not None:
+        window.set_statement_data(df)
+
+    window.show()
+    app.exec()

banking_breakdown/ui/main_window.py

@@ -1,79 +1,15 @@
-import sys
 import typing
 from functools import partial
 
 import pandas as pd
-from PyQt6 import uic, QtGui, QtCore
+from PyQt6 import uic
 from PyQt6.QtCore import Qt, QSortFilterProxyModel
 from PyQt6.QtGui import QPixmap, QAction
 from PyQt6.QtWidgets import QMainWindow, QPushButton, QHBoxLayout, QLabel, \
-    QVBoxLayout, QMenu, QApplication, QTableView, QListView, QInputDialog, \
+    QVBoxLayout, QMenu, QTableView, QInputDialog, QMessageBox, QFileDialog, \
-    QMessageBox, QFileDialog, QListWidget
+    QListWidget
 
-
+from banking_breakdown.ui.pandas_model import PandasModel
-#
-# PandasModel
-#
-
-
-class PandasModel(QtCore.QAbstractTableModel):
-    def __init__(self, df: pd.DataFrame, parent=None):
-        QtCore.QAbstractTableModel.__init__(self, parent)
-        self._data = df
-
-        self._horizontalHeaders = [''] * len(df.columns)
-
-        for i, column in enumerate(df.columns):
-            self.setHeaderData(i, Qt.Orientation.Horizontal, column)
-
-    def setHeaderData(self, section, orientation, data,
-                      role=Qt.ItemDataRole.EditRole):
-
-        if ((orientation == Qt.Orientation.Horizontal)
-                and ((role == Qt.ItemDataRole.DisplayRole)
-                     or (role == Qt.ItemDataRole.EditRole))):
-
-            self._horizontalHeaders[section] = data
-            return True
-        else:
-            return super().setHeaderData(section, orientation, data, role)
-
-    def headerData(self, section, orientation,
-                   role=Qt.ItemDataRole.DisplayRole):
-
-        if (orientation == Qt.Orientation.Horizontal
-                and role == Qt.ItemDataRole.DisplayRole):
-
-            return self._horizontalHeaders[section]
-        else:
-            return super().headerData(section, orientation, role)
-
-    def rowCount(self, parent=None):
-        return len(self._data.values)
-
-    def columnCount(self, parent=None):
-        return self._data.columns.size
-
-    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
-        if not index.isValid():
-            return QtCore.QVariant()
-        if role != Qt.ItemDataRole.DisplayRole:
-            return QtCore.QVariant()
-
-        item = self._data.iloc[index.row()].iloc[index.column()]
-
-        if type(item) is pd.Timestamp:
-            return QtCore.QVariant(item.strftime('%Y-%m-%d'))
-        else:
-            return QtCore.QVariant(str(item))
-
-    def get_dataframe(self):
-        return self._data
-
-
-#
-# MainWindow
-#
 
 
 class MainWindow(QMainWindow):
@@ -306,8 +242,7 @@ class MainWindow(QMainWindow):
         df = self.get_statement_data()
         try:
             df[column_text] \
-                = (pd.to_datetime(df[column_text], format=date_format)
+                = pd.to_datetime(df[column_text], format=date_format)
-                   .dt.strftime('%Y-%m-%d'))
         except:
             QMessageBox.warning(self, "No action performed",
                                 "An error occurred.")
@@ -422,32 +357,3 @@ class MainWindow(QMainWindow):
     def _handle_table_selection_changed(self):
         self._check_enable_clear_button()
         self._check_enable_apply_button()
-
-
-#
-# Other functions
-#
-
-
-def show_main_window(df: pd.DataFrame = None):
-    app = QApplication(sys.argv)
-    window = MainWindow()
-
-    if df is not None:
-        window.set_statement_data(df)
-
-    window.show()
-    app.exec()
-
-
-def main():
-    import os
-    os.chdir("../")
-
-    df = pd.read_csv("res/bank_statement_2023_categorized_renamed.csv")
-
-    show_main_window(df)
-
-
-if __name__ == "__main__":
-    main()

banking_breakdown/ui/pandas_model.py

@@ -0,0 +1,42 @@
+import pandas as pd
+from PyQt6 import QtCore
+from PyQt6.QtCore import Qt
+
+
+class PandasModel(QtCore.QAbstractTableModel):
+    def __init__(self, df: pd.DataFrame, parent=None):
+        QtCore.QAbstractTableModel.__init__(self, parent)
+
+        self._data = df
+        self._horizontalHeaders = list(df.columns)
+
+    # Overloaded functions
+
+    def rowCount(self, parent=None):
+        return len(self._data.values)
+
+    def columnCount(self, parent=None):
+        return self._data.columns.size
+
+    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
+        if not index.isValid():
+            return QtCore.QVariant()
+        if role != Qt.ItemDataRole.DisplayRole:
+            return QtCore.QVariant()
+
+        item = self._data.iloc[index.row(), index.column()]
+        return QtCore.QVariant(str(item))
+
+    def headerData(self, section, orientation,
+                   role=Qt.ItemDataRole.DisplayRole):
+
+        if not ((orientation == Qt.Orientation.Horizontal)
+                and (role == Qt.ItemDataRole.DisplayRole)):
+            return super().headerData(section, orientation, role)
+
+        return self._horizontalHeaders[section]
+
+    # Other functions
+
+    def get_dataframe(self):
+        return self._data