From e958215779bbec11edad343cba4eebb5395a3fc7 Mon Sep 17 00:00:00 2001
From: Andreas Tsouchlos <an.tsouchlos@gmail.com>
Date: Thu, 4 Jan 2024 01:11:21 +0100
Subject: [PATCH] Implement column deletion and renaming

---
 banking_breakdown/ui.py | 57 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 51 insertions(+), 6 deletions(-)

diff --git a/banking_breakdown/ui.py b/banking_breakdown/ui.py
index b0d693a..ead51f1 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,10 @@ class MainWindow(QMainWindow):
         proxyModel.setSourceModel(model)
         self._table_view.setModel(proxyModel)
 
-        self._table_view.resizeColumnsToContents()
+        # self._table_view.resizeColumnsToContents()
+
+    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 +134,49 @@ 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()
+            print(df.columns)
+            df.rename(columns={column_text: new_name}, inplace=True)
+            print(df.columns)
+            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 +221,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)