FEAT: Create a separate class for model toolbox.

Author: genedan

faslr/common/model.py

@@ -47,6 +47,7 @@
     import typing
     from faslr.model import FIBNRModel
     from pandas import DataFrame
+    from PyQt6.QtCore import QAbstractItemModel
     from typing import Any
 
 class FSelectionModel(FAbstractTableModel):
@@ -270,8 +271,6 @@ def select_average_row(self, index: QModelIndex) -> None:
         """
         Makes the ratio selection, i.e., selects the LDFs or loss ratios.
         """
-        for row in self.included_averages_rows:
-            print(row)
         # If double-clicked row is not in the included averages, do nothing.
         if index.row() not in self.included_averages_rows:
             return None
@@ -399,62 +398,77 @@ def __init__(
             self.setWindowTitle(window_title)
 
         self.layout = QVBoxLayout()
-        self.add_average_button = QPushButton("Available Averages")
-        self.add_average_button.setFixedWidth(self.add_average_button.sizeHint().width())
-
-        self.add_average_button.setContentsMargins(
-            2,
-            2,
-            2,
-            2
-        )
 
         # Container widget for upper-right hand tools (add average button, etc.).
-        self.tool_container = QWidget()
-        self.tool_layout = QHBoxLayout()
-        self.tool_container.setLayout(self.tool_layout)
-
-        self.tool_layout.setContentsMargins(
-            0,
-            0,
-            0,
-            0
-        )
-
-        self.tool_layout.addWidget(
-            self.add_average_button
+        self.toolbox = FSelectionModelToolbox(
+            parent=self,
+            averages=averages
         )
 
         self.layout.addWidget(
-            self.tool_container,
+            self.toolbox,
             alignment=Qt.AlignmentFlag.AlignRight
         )
 
+        # If selection model and view has already been subclass, skip, otherwise use base classes.
         if not (hasattr(self, 'selection_model') and hasattr(self, 'selection_model_view')):
             self.selection_model = FSelectionModel(data=data, averages=averages, parent=self)
             self.selection_model_view = FModelView(parent=self)
 
-        self.average_box = FAverageBox(data=averages, parent=self)
-
         self.selection_model_view.setModel(self.selection_model)
 
         self.layout.addWidget(self.selection_model_view)
 
         self.setLayout(self.layout)
 
-        self.add_average_button.clicked.connect(self.open_average_box) # noqa
 
-    def open_average_box(self):
+class FSelectionModelToolbox(QWidget):
+    """
+    Toolbox that appears above the model view in a FSelectionModelWidget. Contains utilities such
+    as a button for adding averages, and a heatmap button.
+    """
+    def __init__(
+            self,
+            parent: FSelectionModelWidget,
+            averages: DataFrame
+    ):
+        super().__init__()
+        self.parent: FSelectionModelWidget = parent
+        self.add_average_button = QPushButton("Available Averages")
+        self.add_average_button.setFixedWidth(self.add_average_button.sizeHint().width())
+
+        self.add_average_button.setContentsMargins(
+            2,
+            2,
+            2,
+            2
+        )
+
+        self.layout = QHBoxLayout()
+        self.setLayout(self.layout)
 
-        self.average_box.show()
+        self.setContentsMargins(
+            0,
+            0,
+            0,
+            0
+        )
 
+        self.average_box = FAverageBox(
+            parent=self,
+            data=averages,
+            selection_model=self.parent.selection_model
+        )
 
-class FSelectionModelToolBox(QWidget):
-    def __init__(self):
-        super().__init__()
+        self.layout.addWidget(self.add_average_button)
 
-        self.layout = QHBoxLayout()
+        self.add_average_button.clicked.connect(self.open_average_box)  # noqa
 
+    def open_average_box(self) -> None:
+        """
+        Opens the dialog box for adding averages to a selection model.
+        """
+        self.average_box.show()
 
 class FModelView(FTableView):
     def __init__(
@@ -470,9 +484,8 @@ def __init__(
     def keyPressEvent(self, e: QKeyEvent) -> None:
         clipboard = QApplication.clipboard()
         if e.matches(QKeySequence.StandardKey.Paste):
-            print('asdfasdf')
             value = clipboard.text()
-            model: FSelectionModel = self.model()
+            model: QAbstractItemModel = self.model()
             for index in self.selectedIndexes():
                 model.setData(index=index, value=value, role=Qt.ItemDataRole.EditRole)
 

faslr/model/average.py

@@ -33,12 +33,13 @@
 
 if TYPE_CHECKING:
     from pandas import DataFrame
-    from faslr.common.model import FSelectionModelWidget
+    from faslr.common.model import FSelectionModel
 
 class FAverageBox(QDialog):
     def __init__(
             self,
-            parent: FSelectionModelWidget = None,
+            parent,
+            selection_model: FSelectionModel = None,
             title: str = None,
             data: DataFrame = None
     ):
@@ -51,6 +52,8 @@ def __init__(
 
         self.parent = parent
 
+        self.selection_model = selection_model
+
         self.model = FAverageModel(
             parent=None,
             data=data,
@@ -85,7 +88,7 @@ def __init__(
         self.setLayout(self.layout)
         self.set_dimensions()
 
-        self.button_box.accepted.connect(self.accept_changes)
+        self.button_box.accepted.connect(self.accept_changes) # noqa
         self.button_box.clicked.connect(self.handle_button_box) # noqa
 
     def set_dimensions(self):
@@ -120,7 +123,7 @@ def accept_changes(self) -> None:
             return
 
         index = QModelIndex()
-        self.parent.selection_model.setData(index=index, value=None)
+        self.selection_model.setData(index=index, value=None)
 
         self.close()