diff --git a/banking_breakdown/ui.py b/banking_breakdown/ui.py index b0d693a..6e815ef 100644 --- a/banking_breakdown/ui.py +++ b/banking_breakdown/ui.py @@ -1,12 +1,14 @@ 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 + QVBoxLayout, QMenu, QApplication, QTableView, QListView, QInputDialog, \ + QMessageBox class PandasModel(QtCore.QAbstractTableModel): @@ -60,6 +62,9 @@ class PandasModel(QtCore.QAbstractTableModel): else: return QtCore.QVariant(str(item)) + def get_dataframe(self): + return self._data + class MainWindow(QMainWindow): def __init__(self): @@ -117,7 +122,16 @@ class MainWindow(QMainWindow): proxyModel.setSourceModel(model) self._table_view.setModel(proxyModel) - self._table_view.resizeColumnsToContents() + if len(df.columns) < 7: + self._table_view.resizeColumnsToContents() + else: + 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) + + 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: @@ -126,11 +140,47 @@ class MainWindow(QMainWindow): def _header_right_clicked(self, pos): context = QMenu(self) - context.addAction(QAction("test 1", self)) - context.addAction(QAction("test 2", self)) - context.addAction(QAction("test 3", self)) + + rename_action = QAction("Rename", self) + delete_action = QAction("Delete", 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)) + + context.addAction(rename_action) + context.addAction(delete_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 _handle_create_click(self): self.add_categories(['asdf']) @@ -175,7 +225,6 @@ def main(): "Transportation", "Social Life", "Other"] df = pd.read_csv("res/bank_statement_2023.csv", delimiter=';') - # df = get_stripped_statement("res/bank_statement_2023.csv") show_main_window(categories, df)