From 14e830ead071e597dd7b4b267e57df97827825c6 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 | 61 +++++++++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 6 deletions(-)

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)