import sys import typing from functools import partial import pandas as pd from PyQt6 import uic, QtGui, QtCore 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, \ QMessageBox, QFileDialog, QListWidget # # 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): 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._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._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) def set_statement_data(self, df: pd.DataFrame): if 'category' not in df.columns: df['category'] = [' '] * len(df.index) model = PandasModel(df) proxyModel = QSortFilterProxyModel(self) proxyModel.setSourceModel(model) self._table_view.setModel(proxyModel) 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) def get_statement_data(self) -> pd.DataFrame: return self._table_view.model().sourceModel().get_dataframe() 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): context = QMenu(self) rename_action = QAction("Rename", self) delete_action = QAction("Delete", self) switch_action = QAction("Switch position", self) column = self._table_view.horizontalHeader().logicalIndexAt(pos) 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.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 _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)) 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 show_main_window(categories: typing.Sequence[str] = None, df: pd.DataFrame = None): app = QApplication(sys.argv) window = MainWindow() if categories is not None: window.add_categories(categories) if df is not None: window.set_statement_data(df) window.show() app.exec() def main(): import os os.chdir("../") categories = ["Food", "Rent & Utilities", "Hobbies", "Education", "Transportation", "Social Life", "Other"] df = pd.read_csv("res/bank_statement_2023_categorized_renamed.csv") show_main_window(categories, df) if __name__ == "__main__": main()