import typing import pandas as pd from PyQt6 import uic from PyQt6.QtCore import Qt, QSortFilterProxyModel from PyQt6.QtGui import QAction from PyQt6.QtWidgets import QMainWindow, QPushButton, QVBoxLayout, \ QTableView, QInputDialog, QMessageBox, QFileDialog, QListWidget from banking_breakdown.ui.pandas_model import PandasModel from banking_breakdown.ui.custom_ui_items import WarningItem, HeaderContextMenu class MainWindow(QMainWindow): def __init__(self): super(MainWindow, self).__init__() uic.loadUi("res/main_window.ui", self) self.setContentsMargins(9, 9, 9, 9) self._warnings = [] # Extract elements 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") # Set up QTableView model self._pandas_model = PandasModel(self) self._proxyModel = QSortFilterProxyModel(self) self._proxyModel.setSourceModel(self._pandas_model) self._table_view.setModel(self._proxyModel) self._proxyModel.setDynamicSortFilter(True) self._proxyModel.setSortRole(Qt.ItemDataRole.EditRole) # Set event handlers 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._handle_header_right_click) self._list_widget.itemSelectionChanged.connect( self._handle_list_selection_changed) self._table_view.selectionModel().selectionChanged.connect( self._handle_table_selection_changed) # # Table data updates # def set_statement_data(self, df: pd.DataFrame): self._pandas_model.set_dataframe(df) self._dataframe_update_callback() self._resize_table_columns_to_content() def get_statement_data(self) -> pd.DataFrame: return self._pandas_model.get_dataframe() def _dataframe_update_callback(self): self._show_warnings() self._update_categories_from_dataframe() def _resize_table_columns_to_content(self): df = self._pandas_model.get_dataframe() # Resize columns 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) # # List data updates # def _add_categories(self, categories: typing.Sequence[str]): for category in categories: self._list_widget.addItem(category) def _update_categories_from_dataframe(self): df = self._pandas_model.get_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 != ' ']) # # Warnings # def _add_warning_item(self, text: str): warning_item = WarningItem(text=text, parent=self) self._warning_layout.addLayout(warning_item) self._warnings.append(warning_item) def _show_warnings(self): for warning_item in self._warnings: warning_item.hide() self._warning_layout.removeItem(warning_item) df = self._pandas_model.get_dataframe() if 't' not in df.columns: self._add_warning_item( "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_item( "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_item( "The column 'balance' does not exist. Please rename the column" " containing the balance after each transaction to 'balance'") # # Event handlers # def _handle_header_right_click(self, pos): column = self._table_view.horizontalHeader().logicalIndexAt(pos) context = HeaderContextMenu(parent=self, column=column, pandas_model=self._pandas_model, callback=self._dataframe_update_callback) context.exec(self.sender().mapToGlobal(pos)) 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._pandas_model.get_dataframe() df.loc[row_indices, 'category'] = category self._pandas_model.set_dataframe(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) # # Enable / Disable buttons # 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()