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._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) .dt.strftime('%Y-%m-%d')) 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() # # 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()