Separate WarningItem and HeaderContextMenu from MainWindow

banking_breakdown/ui/main_window.py

@@ -12,12 +12,183 @@ from PyQt6.QtWidgets import QMainWindow, QPushButton, QHBoxLayout, QLabel, \
 from banking_breakdown.ui.pandas_model import PandasModel
 
 
+class WarningItem(QHBoxLayout):
+    def __init__(self, text: str, parent=None):
+        super(WarningItem, self).__init__()
+
+        self._warningIcon = QLabel()
+        pixmap = QPixmap("res/warning.png")
+        self._warningIcon.setPixmap(pixmap)
+
+        self._label = QLabel(text)
+        self._label.setWordWrap(True)
+
+        self.addWidget(self._warningIcon)
+        self.addWidget(self._label)
+
+        self.setStretch(0, 0)
+        self.setStretch(1, 1)
+
+    def hide(self):
+        self._label.hide()
+        self._warningIcon.hide()
+
+
+class HeaderContextMenu(QMenu):
+    def __init__(self, column, pandas_model: PandasModel, callback=None,
+                 parent=None):
+        super(HeaderContextMenu, self).__init__()
+
+        self._column = column
+        self._pandas_model = pandas_model
+        self._callback = callback
+
+        self._column_text \
+            = self._pandas_model.headerData(self._column,
+                                            Qt.Orientation.Horizontal)
+
+        # Define assign action
+
+        assign_menu = QMenu("Assign type", self)
+        assign_date_action = QAction("date", self)
+        assign_float_action = QAction("float", self)
+
+        assign_menu.addAction(assign_date_action)
+        assign_menu.addAction(assign_float_action)
+
+        assign_date_action.triggered.connect(self._assign_date_handler)
+        assign_float_action.triggered.connect(self._assign_float_handler)
+
+        # Define other actions
+
+        rename_action = QAction("Rename", self)
+        delete_action = QAction("Delete", self)
+        switch_action = QAction("Switch position", self)
+
+        rename_action.triggered.connect(self._rename_handler)
+        delete_action.triggered.connect(self._delete_handler)
+        switch_action.triggered.connect(self._switch_handler)
+
+        # Add actions to menu
+
+        self.addAction(rename_action)
+        self.addAction(delete_action)
+        self.addAction(switch_action)
+        self.addAction(assign_menu.menuAction())
+
+    def _rename_handler(self):
+        new_name, flag = QInputDialog.getText(self, "Rename column",
+                                              "New name:",
+                                              text=self._column_text)
+
+        if not flag:
+            return
+
+        if (new_name != self._column_text) and (new_name != ''):
+            df = self._pandas_model.get_dataframe()
+            df = df.rename(columns={self._column_text: new_name})
+            self._pandas_model.set_dataframe(df)
+
+            if self._callback:
+                self._callback()
+
+    def _delete_handler(self):
+        button = QMessageBox.question(self, "Delete column",
+                                      f"Are you sure you want to delete"
+                                      f" column '{self._column_text}'?")
+
+        if button == QMessageBox.StandardButton.Yes:
+            df = self._pandas_model.get_dataframe()
+            df = df.iloc[:, [j for j, c
+                             in enumerate(df.columns) if j != self._column]]
+            self._pandas_model.set_dataframe(df)
+
+            if self._callback:
+                self._callback()
+
+    def _switch_handler(self):
+        df = self._pandas_model.get_dataframe()
+        columns = [column for column in df.columns
+                   if column != self._column_text]
+
+        other_name, flag = QInputDialog.getItem(self, "Switch column position",
+                                                f"Switch position of colum"
+                                                f" '{self._column_text}' with:",
+                                                columns, editable=False)
+
+        if not flag:
+            return
+
+        column_titles = list(df.columns)
+        index1, index2 = column_titles.index(
+            self._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._pandas_model.set_dataframe(df)
+
+        if self._callback:
+            self._callback()
+
+    def _assign_date_handler(self):
+        date_format, flag = QInputDialog.getText(self, "Format",
+                                                 "Format:",
+                                                 text="%d.%m.%Y")
+
+        if not flag:
+            return
+
+        df = self._pandas_model.get_dataframe()
+        try:
+            df[self._column_text] \
+                = pd.to_datetime(df[self._column_text], format=date_format)
+        except:
+            QMessageBox.warning(self, "No action performed",
+                                "An error occurred.")
+        self._pandas_model.set_dataframe(df)
+
+        if self._callback:
+            self._callback()
+
+    def _assign_float_handler(self):
+        chars = ['.', ',']
+        decimal_sep, flag = QInputDialog.getItem(self, "Decimal separator",
+                                                 "Decimal separator:",
+                                                 chars, editable=False)
+
+        if not flag:
+            return
+
+        df = self._pandas_model.get_dataframe()
+
+        try:
+            if decimal_sep == ',':
+                df[self._column_text] \
+                    = df[self._column_text].str.replace(',', '.').astype(float)
+            else:
+                df[self._column_text] = df[self._column_text].astype(float)
+        except:
+            QMessageBox.warning(self, "No action performed",
+                                "An error occurred.")
+
+        self._pandas_model.set_dataframe(df)
+
+        if self._callback:
+            self._callback()
+
+
 class MainWindow(QMainWindow):
     def __init__(self):
         super(MainWindow, self).__init__()
+
         uic.loadUi("res/main_window.ui", self)
         self.setContentsMargins(9, 9, 9, 9)
 
+        self._warnings = []
+
+        # Extract elements
+
         self._warning_layout \
             = self.findChild(QVBoxLayout, "warningLayout")
         self._create_button \
@@ -35,7 +206,15 @@ class MainWindow(QMainWindow):
         self._action_save \
             = self.findChild(QAction, "actionSave")
 
-        self._warnings = []
+        # Set up QTableView model
+
+        self._pandas_model = PandasModel(self)
+        self._proxyModel = QSortFilterProxyModel(self)
+        self._proxyModel.setSourceModel(self._pandas_model)
+        self._table_view.setModel(self._proxyModel)
+        self._proxyModel.setDynamicSortFilter(False)
+
+        # Set event handlers
 
         self._create_button.clicked.connect(self._handle_create_click)
         self._delete_button.clicked.connect(self._handle_delete_click)
@@ -52,17 +231,23 @@ class MainWindow(QMainWindow):
         self._list_widget.itemSelectionChanged.connect(
             self._handle_list_selection_changed)
 
-        self._pandas_model = PandasModel(self)
+        self._table_view.selectionModel().selectionChanged.connect(
-        self._proxyModel = QSortFilterProxyModel(self)
+            self._handle_table_selection_changed)
-        self._proxyModel.setSourceModel(self._pandas_model)
-        self._table_view.setModel(self._proxyModel)
-        self._proxyModel.setDynamicSortFilter(False)
 
     def set_statement_data(self, df: pd.DataFrame):
-        if 'category' not in df.columns:
-            df['category'] = [' '] * len(df.index)
-
         self._pandas_model.set_dataframe(df)
+        self._dataframe_update_callback()
+
+    def get_statement_data(self) -> pd.DataFrame:
+        return self._pandas_model.get_dataframe()
+
+    def _dataframe_update_callback(self):
+        df = self._pandas_model.get_dataframe()
+
+        self._show_warnings()
+        self._update_categories_from_dataframe(df)
+
+        # Resize columns
 
         if len(df.columns) < 10:  # Experimentally determined threshold
             # Properly resize columns (takes longer)
@@ -74,18 +259,9 @@ class MainWindow(QMainWindow):
                                len(col))
                 self._table_view.setColumnWidth(i, max_char * 10)
 
-        self._show_warnings(df)
+    def _add_categories(self, categories: typing.Sequence[str]):
-        self._update_categories_from_dataframe(df)
+        for category in categories:
-
+            self._list_widget.addItem(category)
-        # Upon reloading the data, the selection is cleared but no signal
-        # is generated. Manually call the signal handler
-        self._handle_table_selection_changed()
-
-        self._table_view.selectionModel().selectionChanged.connect(
-            self._handle_table_selection_changed)
-
-    def get_statement_data(self) -> pd.DataFrame:
-        return self._table_view.model().sourceModel().get_dataframe()
 
     def _update_categories_from_dataframe(self, df: pd.DataFrame):
         df_categories = df['category'].unique()
@@ -96,186 +272,42 @@ class MainWindow(QMainWindow):
         self._add_categories([category for category
                               in missing if category != ' '])
 
-    def _add_categories(self, categories: typing.Sequence[str]):
+    def _add_warning_item(self, text: str):
-        for category in categories:
+        warning_item = WarningItem(text=text, parent=self)
-            self._list_widget.addItem(category)
 
-    def _add_warning_text(self, text: str):
+        self._warning_layout.addLayout(warning_item)
-        layout = QHBoxLayout()
+        self._warnings.append(warning_item)
 
-        warningIcon = QLabel()
+    def _show_warnings(self):
-        pixmap = QPixmap("res/warning.png")
+        for warning_item in self._warnings:
-        warningIcon.setPixmap(pixmap)
+            warning_item.hide()
+            self._warning_layout.removeItem(warning_item)
 
-        label = QLabel(text)
+        df = self._pandas_model.get_dataframe()
-        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(
+            self._add_warning_item(
-                "The column 't' does not exist. Please rename the"
+                "The column 't' does not exist. Please rename the column"
-                " column containing the dates of the transactions"
+                " containing the dates of the transactions to 't'.")
-                " to 't'.")
 
         if 'value' not in df.columns:
-            self._add_warning_text(
+            self._add_warning_item(
-                "The column 'value' does not exist. Please rename"
+                "The column 'value' does not exist. Please rename the column"
-                " the column containing the values of the"
+                " containing the values of the transactions to 'value'.")
-                " transactions to 'value'.")
 
         if 'balance' not in df.columns:
-            self._add_warning_text(
+            self._add_warning_item(
-                "The column 'balance' does not exist. Please"
+                "The column 'balance' does not exist. Please rename the column"
-                " rename the column containing the balance"
+                " containing the balance after each transaction to 'balance'")
-                " after each transaction to 'balance'")
 
     def _header_right_clicked(self, pos):
         column = self._table_view.horizontalHeader().logicalIndexAt(pos)
 
-        context = QMenu(self)
+        context = HeaderContextMenu(parent=self, column=column,
-
+                                    pandas_model=self._pandas_model,
-        rename_action = QAction("Rename", self)
+                                    callback=self._dataframe_update_callback)
-        delete_action = QAction("Delete", self)
-        switch_action = QAction("Switch position", self)
-
-        assign_menu = QMenu("Assign type", self)
-        date_action = QAction("date", self)
-        float_action = QAction("float", self)
-        assign_menu.addAction(date_action)
-        assign_menu.addAction(float_action)
-        date_action.triggered.connect(
-            partial(self._header_assign_date_handler, column))
-        float_action.triggered.connect(
-            partial(self._header_assign_float_handler, column))
-
-        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.addAction(assign_menu.menuAction())
-
         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 _header_assign_date_handler(self, column):
-        model = self._table_view.horizontalHeader().model()
-        column_text = model.headerData(column, Qt.Orientation.Horizontal)
-
-        date_format, flag = QInputDialog.getText(self, "Format",
-                                                 "Format:",
-                                                 text="%d.%m.%Y")
-
-        if not flag:
-            return
-
-        df = self.get_statement_data()
-        try:
-            df[column_text] \
-                = pd.to_datetime(df[column_text], format=date_format)
-        except:
-            QMessageBox.warning(self, "No action performed",
-                                "An error occurred.")
-        self.set_statement_data(df)
-
-    def _header_assign_float_handler(self, column):
-        model = self._table_view.horizontalHeader().model()
-        column_text = model.headerData(column, Qt.Orientation.Horizontal)
-
-        chars = ['.', ',']
-        decimal_sep, flag = QInputDialog.getItem(self, "Decimal separator",
-                                                 "Decimal separator:",
-                                                 chars, editable=False)
-
-        if not flag:
-            return
-
-        df = self.get_statement_data()
-
-        try:
-            if decimal_sep == ',':
-                df[column_text] \
-                    = df[column_text].str.replace(',', '.').astype(float)
-            else:
-                df[column_text] = df[column_text].astype(float)
-        except:
-            QMessageBox.warning(self, "No action performed",
-                                "An error occurred.")
-
-        self.set_statement_data(df)
-
     def _handle_create_click(self):
         new_name, flag = QInputDialog.getText(self, "Create category",
                                               "New category:",