From 393f654a573ee683dfd52fe5b8836be989c946df Mon Sep 17 00:00:00 2001
From: Andreas Tsouchlos <an.tsouchlos@gmail.com>
Date: Thu, 4 Jan 2024 18:36:13 +0100
Subject: [PATCH] Change the way PandasModel operates

---
 banking_breakdown/ui/main_window.py  |  6 +++--
 banking_breakdown/ui/pandas_model.py | 33 ++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/banking_breakdown/ui/main_window.py b/banking_breakdown/ui/main_window.py
index 27dab90..8eb7848 100644
--- a/banking_breakdown/ui/main_window.py
+++ b/banking_breakdown/ui/main_window.py
@@ -52,15 +52,17 @@ class MainWindow(QMainWindow):
         self._list_widget.itemSelectionChanged.connect(
             self._handle_list_selection_changed)
 
+        self._pandas_model = PandasModel(self)
         self._proxyModel = QSortFilterProxyModel(self)
+        self._proxyModel.setSourceModel(self._pandas_model)
         self._table_view.setModel(self._proxyModel)
+        self._proxyModel.setDynamicSortFilter(False)
 
     def set_statement_data(self, df: pd.DataFrame):
         if 'category' not in df.columns:
             df['category'] = [' '] * len(df.index)
 
-        model = PandasModel(df)
-        self._proxyModel.setSourceModel(model)
+        self._pandas_model.set_dataframe(df)
 
         if len(df.columns) < 10:  # Experimentally determined threshold
             # Properly resize columns (takes longer)
diff --git a/banking_breakdown/ui/pandas_model.py b/banking_breakdown/ui/pandas_model.py
index 53005e3..fbd2069 100644
--- a/banking_breakdown/ui/pandas_model.py
+++ b/banking_breakdown/ui/pandas_model.py
@@ -3,20 +3,31 @@ from PyQt6 import QtCore
 from PyQt6.QtCore import Qt
 
 
+def _get_str_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+    """Return a given dataframe with all values turned into strings.
+
+    When the data given to the PandasModel class contains non-strings,
+    an attached QTableView seems to respond rather slowly. This function
+    turns all data in the DataFrame into strings, yielding a better experience.
+    """
+    return df.astype(str)
+
+
 class PandasModel(QtCore.QAbstractTableModel):
-    def __init__(self, df: pd.DataFrame, parent=None):
+    def __init__(self, parent=None):
         QtCore.QAbstractTableModel.__init__(self, parent)
 
-        self._data = df
-        self._horizontalHeaders = list(df.columns)
+        self._data = pd.DataFrame()
+        self._data_str = pd.DataFrame()
+        self._horizontalHeaders = None
 
     # Overloaded functions
 
     def rowCount(self, parent=None):
-        return len(self._data.values)
+        return len(self._data_str.values)
 
     def columnCount(self, parent=None):
-        return self._data.columns.size
+        return self._data_str.columns.size
 
     def data(self, index, role=Qt.ItemDataRole.DisplayRole):
         if not index.isValid():
@@ -24,7 +35,7 @@ class PandasModel(QtCore.QAbstractTableModel):
         if role != Qt.ItemDataRole.DisplayRole:
             return QtCore.QVariant()
 
-        item = self._data.iloc[index.row(), index.column()]
+        item = self._data_str.iloc[index.row(), index.column()]
         return QtCore.QVariant(str(item))
 
     def headerData(self, section, orientation,
@@ -38,5 +49,13 @@ class PandasModel(QtCore.QAbstractTableModel):
 
     # Other functions
 
-    def get_dataframe(self):
+    def set_dataframe(self, df):
+        self._data = df
+        self._data_str = _get_str_dataframe(df)
+        self._horizontalHeaders = list(df.columns)
+
+        self.layoutAboutToBeChanged.emit()
+        self.layoutChanged.emit()
+
+    def get_dataframe(self) -> pd.DataFrame:
         return self._data