Implement column deletion and renaming

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)
 
+        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))
+        rename_action = QAction("Rename", self)
-        context.addAction(QAction("test 3", 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)