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

@@ -0,0 +1,359 @@
+import typing
+from functools import partial
+
+import pandas as pd
+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, QTableView, QInputDialog, QMessageBox, QFileDialog, \
+    QListWidget
+
+from banking_breakdown.ui.pandas_model import PandasModel
+
+
+class MainWindow(QMainWindow):
+    def __init__(self):
+        super(MainWindow, self).__init__()
+        uic.loadUi("res/main_window.ui", self)
+        self.setContentsMargins(9, 9, 9, 9)
+
+        self._warning_layout \
+            = self.findChild(QVBoxLayout, "warningLayout")
+        self._create_button \
+            = self.findChild(QPushButton, "createCategoryButton")
+        self._delete_button \
+            = self.findChild(QPushButton, "deleteCategoryButton")
+        self._apply_button \
+            = self.findChild(QPushButton, "applyCategoryButton")
+        self._clear_button \
+            = self.findChild(QPushButton, "clearCategoryButton")
+        self._list_widget \
+            = self.findChild(QListWidget, "categoryListWidget")
+        self._table_view \
+            = self.findChild(QTableView, "transactionTableView")
+        self._action_save \
+            = self.findChild(QAction, "actionSave")
+
+        self._warnings = []
+
+        self._create_button.clicked.connect(self._handle_create_click)
+        self._delete_button.clicked.connect(self._handle_delete_click)
+        self._clear_button.clicked.connect(self._handle_clear_click)
+        self._apply_button.clicked.connect(self._handle_apply_click)
+        self._list_widget.itemActivated.connect(self._handle_item_double_click)
+
+        self._action_save.triggered.connect(self._handle_save)
+
+        header = self._table_view.horizontalHeader()
+        header.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
+        header.customContextMenuRequested.connect(self._header_right_clicked)
+
+        self._list_widget.itemSelectionChanged.connect(
+            self._handle_list_selection_changed)
+
+        self._proxyModel = QSortFilterProxyModel(self)
+        self._table_view.setModel(self._proxyModel)
+
+    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)
+
+        if len(df.columns) < 10:  # Experimentally determined threshold
+            # Properly resize columns (takes longer)
+            self._table_view.resizeColumnsToContents()
+        else:
+            # Quickly approximate sizes
+            for i, col in enumerate(df.columns):
+                max_char = max(max([len(str(x)) for x in df[col].values]),
+                               len(col))
+                self._table_view.setColumnWidth(i, max_char * 10)
+
+        self._show_warnings(df)
+        self._update_categories_from_dataframe(df)
+
+        # Upon reloading the data, the selection is cleared but no signal
+        # is generated. Manually call the signal handler
+        self._handle_table_selection_changed()
+
+        self._table_view.selectionModel().selectionChanged.connect(
+            self._handle_table_selection_changed)
+
+    def get_statement_data(self) -> pd.DataFrame:
+        return self._table_view.model().sourceModel().get_dataframe()
+
+    def _update_categories_from_dataframe(self, df: pd.DataFrame):
+        df_categories = df['category'].unique()
+        current_categories = [self._list_widget.item(x).text() for x
+                              in range(self._list_widget.count())]
+
+        missing = list(set(df_categories) - set(current_categories))
+        self._add_categories([category for category
+                              in missing if category != ' '])
+
+    def _add_categories(self, categories: typing.Sequence[str]):
+        for category in categories:
+            self._list_widget.addItem(category)
+
+    def _add_warning_text(self, text: str):
+        layout = QHBoxLayout()
+
+        warningIcon = QLabel()
+        pixmap = QPixmap("res/warning.png")
+        warningIcon.setPixmap(pixmap)
+
+        label = QLabel(text)
+        label.setWordWrap(True)
+
+        layout.addWidget(warningIcon)
+        layout.addWidget(label)
+
+        layout.setStretch(0, 0)
+        layout.setStretch(1, 1)
+
+        self._warning_layout.addLayout(layout)
+        self._warnings.append((layout, warningIcon, label))
+
+    def _show_warnings(self, df: pd.DataFrame):
+        for (layout, icon, text) in self._warnings:
+            icon.hide()
+            text.hide()
+            self._warning_layout.removeItem(layout)
+
+        if 't' not in df.columns:
+            self._add_warning_text(
+                "The column 't' does not exist. Please rename the"
+                " column containing the dates of the transactions"
+                " to 't'.")
+
+        if 'value' not in df.columns:
+            self._add_warning_text(
+                "The column 'value' does not exist. Please rename"
+                " the column containing the values of the"
+                " transactions to 'value'.")
+
+        if 'balance' not in df.columns:
+            self._add_warning_text(
+                "The column 'balance' does not exist. Please"
+                " rename the column containing the balance"
+                " after each transaction to 'balance'")
+
+    def _header_right_clicked(self, pos):
+        column = self._table_view.horizontalHeader().logicalIndexAt(pos)
+
+        context = QMenu(self)
+
+        rename_action = QAction("Rename", self)
+        delete_action = QAction("Delete", self)
+        switch_action = QAction("Switch position", self)
+
+        assign_menu = QMenu("Assign type", self)
+        date_action = QAction("date", self)
+        float_action = QAction("float", self)
+        assign_menu.addAction(date_action)
+        assign_menu.addAction(float_action)
+        date_action.triggered.connect(
+            partial(self._header_assign_date_handler, column))
+        float_action.triggered.connect(
+            partial(self._header_assign_float_handler, column))
+
+        rename_action.triggered.connect(
+            partial(self._header_rename_handler, column))
+        delete_action.triggered.connect(
+            partial(self._header_delete_handler, column))
+        switch_action.triggered.connect(
+            partial(self._header_switch_handler, column))
+
+        context.addAction(rename_action)
+        context.addAction(delete_action)
+        context.addAction(switch_action)
+        context.addAction(assign_menu.menuAction())
+
+        context.exec(self.sender().mapToGlobal(pos))
+
+    def _header_rename_handler(self, column):
+        model = self._table_view.horizontalHeader().model()
+        column_text = model.headerData(column, Qt.Orientation.Horizontal)
+
+        new_name, flag = QInputDialog.getText(self, "Rename column",
+                                              "New name:",
+                                              text=column_text)
+
+        if not flag:
+            return
+
+        if (new_name != column_text) and (new_name != ''):
+            df = self.get_statement_data()
+            df = df.rename(columns={column_text: new_name})
+            self.set_statement_data(df)
+
+    def _header_delete_handler(self, column):
+        model = self._table_view.horizontalHeader().model()
+        column_text = model.headerData(column, Qt.Orientation.Horizontal)
+
+        button = QMessageBox.question(self, "Delete column",
+                                      f"Are you sure you want to delete"
+                                      f" column '{column_text}'?")
+
+        if button == QMessageBox.StandardButton.Yes:
+            df = self.get_statement_data()
+            df = df.iloc[:, [j for j, c
+                             in enumerate(df.columns) if j != column]]
+            self.set_statement_data(df)
+
+    def _header_switch_handler(self, column):
+        model = self._table_view.horizontalHeader().model()
+        column_text = model.headerData(column, Qt.Orientation.Horizontal)
+
+        df = self.get_statement_data()
+        columns = [column for column in df.columns if column != column_text]
+
+        other_name, flag = QInputDialog.getItem(self, "Switch column position",
+                                                f"Switch position of colum"
+                                                f" '{column_text}' with:",
+                                                columns, editable=False)
+
+        if not flag:
+            return
+
+        column_titles = list(df.columns)
+        index1, index2 = column_titles.index(column_text), column_titles.index(
+            other_name)
+        column_titles[index1], column_titles[index2] \
+            = column_titles[index2], column_titles[index1]
+
+        df = df.reindex(columns=column_titles)
+        self.set_statement_data(df)
+
+    def _header_assign_date_handler(self, column):
+        model = self._table_view.horizontalHeader().model()
+        column_text = model.headerData(column, Qt.Orientation.Horizontal)
+
+        date_format, flag = QInputDialog.getText(self, "Format",
+                                                 "Format:",
+                                                 text="%d.%m.%Y")
+
+        if not flag:
+            return
+
+        df = self.get_statement_data()
+        try:
+            df[column_text] \
+                = pd.to_datetime(df[column_text], format=date_format)
+        except:
+            QMessageBox.warning(self, "No action performed",
+                                "An error occurred.")
+        self.set_statement_data(df)
+
+    def _header_assign_float_handler(self, column):
+        model = self._table_view.horizontalHeader().model()
+        column_text = model.headerData(column, Qt.Orientation.Horizontal)
+
+        chars = ['.', ',']
+        decimal_sep, flag = QInputDialog.getItem(self, "Decimal separator",
+                                                 "Decimal separator:",
+                                                 chars, editable=False)
+
+        if not flag:
+            return
+
+        df = self.get_statement_data()
+
+        try:
+            if decimal_sep == ',':
+                df[column_text] \
+                    = df[column_text].str.replace(',', '.').astype(float)
+            else:
+                df[column_text] = df[column_text].astype(float)
+        except:
+            QMessageBox.warning(self, "No action performed",
+                                "An error occurred.")
+
+        self.set_statement_data(df)
+
+    def _handle_create_click(self):
+        new_name, flag = QInputDialog.getText(self, "Create category",
+                                              "New category:",
+                                              text="Category")
+        if not flag:
+            return
+
+        current_items = [self._list_widget.item(x).text() for x
+                         in range(self._list_widget.count())]
+
+        if new_name not in current_items:
+            self._add_categories([new_name])
+        else:
+            QMessageBox.warning(self, "No action performed",
+                                f"Category '{new_name}' already exists.")
+
+    def _handle_delete_click(self):
+        selected_item = self._list_widget.selectedItems()[0]
+
+        button = QMessageBox.question(self, "Delete category",
+                                      f"Are you sure you want to delete"
+                                      f" category '{selected_item.text()}'?")
+
+        if button == QMessageBox.StandardButton.Yes:
+            self._list_widget.takeItem(self._list_widget.row(selected_item))
+            df = self.get_statement_data()
+            df.loc[df['category'] == selected_item.text(), 'category'] = ' '
+
+    def _handle_clear_click(self):
+        self._assign_category(' ')
+
+    def _assign_category(self, category: str):
+        indexes = self._table_view.selectionModel().selectedRows()
+
+        row_indices = [self._table_view.model().mapToSource(index).row()
+                       for index in indexes]
+
+        df = self.get_statement_data()
+        df.loc[row_indices, 'category'] = category
+        self.set_statement_data(df)
+
+    def _handle_apply_click(self):
+        category = self._list_widget.selectedItems()[0].text()
+        self._assign_category(category)
+
+    def _handle_item_double_click(self, item):
+        self._assign_category(item.text())
+
+    def _handle_save(self):
+        filename, _ = QFileDialog.getSaveFileName(self, 'Save File')
+
+        if filename == '':
+            return
+
+        df = self.get_statement_data()
+        df.to_csv(filename, index=False)
+
+    def _check_enable_delete_button(self):
+        if len(self._list_widget.selectedItems()) > 0:
+            self._delete_button.setEnabled(True)
+        else:
+            self._delete_button.setEnabled(False)
+
+    def _check_enable_clear_button(self):
+        if len(self._table_view.selectionModel().selectedRows()) > 0:
+            self._clear_button.setEnabled(True)
+        else:
+            self._clear_button.setEnabled(False)
+
+    def _check_enable_apply_button(self):
+        if ((len(self._table_view.selectionModel().selectedRows()) > 0)
+                and (len(self._list_widget.selectedItems()) > 0)):
+            self._apply_button.setEnabled(True)
+        else:
+            self._apply_button.setEnabled(False)
+
+    def _handle_list_selection_changed(self):
+        self._check_enable_delete_button()
+        self._check_enable_apply_button()
+
+    def _handle_table_selection_changed(self):
+        self._check_enable_clear_button()
+        self._check_enable_apply_button()

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