341 lines
12 KiB
Python
341 lines
12 KiB
Python
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, QInputDialog, \
|
|
QMessageBox, QFileDialog, QListWidget
|
|
|
|
|
|
#
|
|
# PandasModel
|
|
#
|
|
|
|
|
|
class PandasModel(QtCore.QAbstractTableModel):
|
|
def __init__(self, df: pd.DataFrame, parent=None):
|
|
QtCore.QAbstractTableModel.__init__(self, parent)
|
|
self._data = df
|
|
|
|
self._horizontalHeaders = [''] * len(df.columns)
|
|
|
|
for i, column in enumerate(df.columns):
|
|
self.setHeaderData(i, Qt.Orientation.Horizontal, column)
|
|
|
|
def setHeaderData(self, section, orientation, data,
|
|
role=Qt.ItemDataRole.EditRole):
|
|
|
|
if ((orientation == Qt.Orientation.Horizontal)
|
|
and ((role == Qt.ItemDataRole.DisplayRole)
|
|
or (role == Qt.ItemDataRole.EditRole))):
|
|
|
|
self._horizontalHeaders[section] = data
|
|
return True
|
|
else:
|
|
return super().setHeaderData(section, orientation, data, role)
|
|
|
|
def headerData(self, section, orientation,
|
|
role=Qt.ItemDataRole.DisplayRole):
|
|
|
|
if (orientation == Qt.Orientation.Horizontal
|
|
and role == Qt.ItemDataRole.DisplayRole):
|
|
|
|
return self._horizontalHeaders[section]
|
|
else:
|
|
return super().headerData(section, orientation, role)
|
|
|
|
def rowCount(self, parent=None):
|
|
return len(self._data.values)
|
|
|
|
def columnCount(self, parent=None):
|
|
return self._data.columns.size
|
|
|
|
def data(self, index, role=Qt.ItemDataRole.DisplayRole):
|
|
if not index.isValid():
|
|
return QtCore.QVariant()
|
|
if role != Qt.ItemDataRole.DisplayRole:
|
|
return QtCore.QVariant()
|
|
|
|
item = self._data.iloc[index.row()].iloc[index.column()]
|
|
|
|
if type(item) is pd.Timestamp:
|
|
return QtCore.QVariant(item.strftime('%Y-%m-%d'))
|
|
else:
|
|
return QtCore.QVariant(str(item))
|
|
|
|
def get_dataframe(self):
|
|
return self._data
|
|
|
|
|
|
#
|
|
# MainWindow
|
|
#
|
|
|
|
|
|
class MainWindow(QMainWindow):
|
|
def __init__(self):
|
|
super(MainWindow, self).__init__()
|
|
uic.loadUi("res/main_window.ui", self)
|
|
self.setContentsMargins(9, 9, 9, 9)
|
|
|
|
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._list_widget \
|
|
= self.findChild(QListWidget, "categoryListWidget")
|
|
self._table_view \
|
|
= self.findChild(QTableView, "transactionTableView")
|
|
self._action_save \
|
|
= self.findChild(QAction, "actionSave")
|
|
|
|
self._warnings = []
|
|
|
|
self._create_button.clicked.connect(self._handle_create_click)
|
|
self._delete_button.clicked.connect(self._handle_delete_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._header_right_clicked)
|
|
|
|
def set_statement_data(self, df: pd.DataFrame):
|
|
if 'category' not in df.columns:
|
|
df['category'] = [' '] * len(df.index)
|
|
|
|
model = PandasModel(df)
|
|
proxyModel = QSortFilterProxyModel(self)
|
|
proxyModel.setSourceModel(model)
|
|
self._table_view.setModel(proxyModel)
|
|
|
|
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)
|
|
|
|
self._show_warnings(df)
|
|
|
|
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:
|
|
self._list_widget.addItem(category)
|
|
|
|
def _add_warning_text(self, text: str):
|
|
layout = QHBoxLayout()
|
|
|
|
warningIcon = QLabel()
|
|
pixmap = QPixmap("res/warning.png")
|
|
warningIcon.setPixmap(pixmap)
|
|
|
|
label = QLabel(text)
|
|
label.setWordWrap(True)
|
|
|
|
layout.addWidget(warningIcon)
|
|
layout.addWidget(label)
|
|
|
|
layout.setStretch(0, 0)
|
|
layout.setStretch(1, 1)
|
|
|
|
self._warning_layout.addLayout(layout)
|
|
self._warnings.append((layout, warningIcon, label))
|
|
|
|
def _show_warnings(self, df: pd.DataFrame):
|
|
for (layout, icon, text) in self._warnings:
|
|
icon.hide()
|
|
text.hide()
|
|
self._warning_layout.removeItem(layout)
|
|
|
|
if 't' not in df.columns:
|
|
self._add_warning_text(
|
|
"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_text(
|
|
"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_text(
|
|
"The column 'balance' does not exist. Please"
|
|
" rename the column containing the balance"
|
|
" after each transaction to 'balance'")
|
|
|
|
def _header_right_clicked(self, pos):
|
|
context = QMenu(self)
|
|
|
|
rename_action = QAction("Rename", self)
|
|
delete_action = QAction("Delete", self)
|
|
switch_action = QAction("Switch position", 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))
|
|
switch_action.triggered.connect(
|
|
partial(self._header_switch_handler, column))
|
|
|
|
context.addAction(rename_action)
|
|
context.addAction(delete_action)
|
|
context.addAction(switch_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, flag = QInputDialog.getText(self, "Rename column",
|
|
"New name:",
|
|
text=column_text)
|
|
|
|
if not flag:
|
|
return
|
|
|
|
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 _header_switch_handler(self, column):
|
|
model = self._table_view.horizontalHeader().model()
|
|
column_text = model.headerData(column, Qt.Orientation.Horizontal)
|
|
|
|
df = self.get_statement_data()
|
|
columns = [column for column in df.columns if column != column_text]
|
|
|
|
other_name, flag = QInputDialog.getItem(self, "Switch column position",
|
|
f"Switch position of colum"
|
|
f" '{column_text}' with:",
|
|
columns, editable=False)
|
|
|
|
if not flag:
|
|
return
|
|
|
|
column_titles = list(df.columns)
|
|
index1, index2 = column_titles.index(column_text), column_titles.index(
|
|
other_name)
|
|
column_titles[index1], column_titles[index2] \
|
|
= column_titles[index2], column_titles[index1]
|
|
|
|
df = df.reindex(columns=column_titles)
|
|
self.set_statement_data(df)
|
|
|
|
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))
|
|
|
|
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.get_statement_data()
|
|
df.loc[row_indices, 'category'] = category
|
|
self.set_statement_data(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)
|
|
|
|
|
|
def show_main_window(categories: typing.Sequence[str] = None,
|
|
df: pd.DataFrame = None):
|
|
app = QApplication(sys.argv)
|
|
window = MainWindow()
|
|
|
|
if categories is not None:
|
|
window.add_categories(categories)
|
|
|
|
if df is not None:
|
|
window.set_statement_data(df)
|
|
|
|
window.show()
|
|
app.exec()
|
|
|
|
|
|
def main():
|
|
import os
|
|
os.chdir("../")
|
|
|
|
categories = ["Food", "Rent & Utilities", "Hobbies", "Education",
|
|
"Transportation", "Social Life", "Other"]
|
|
|
|
df = pd.read_csv("res/bank_statement_2023_categorized_renamed.csv")
|
|
|
|
show_main_window(categories, df)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|