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 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 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(QListView, "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._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 _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 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 _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 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 _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, _ = QInputDialog.getText(self, "Rename column", "New name:", text=column_text) 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() other_name, _ = QInputDialog.getItem(self, "Switch column position", f"Switch position of colum" f" '{column_text}' with:", df.columns, editable=False) 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): self.add_categories(['asdf']) def _handle_delete_click(self): print(self._list_widget.selectedIndexes()) def _handle_apply_click(self): pass 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()