Insert parameter after selected row

Author: StephenNneji

rascal2/widgets/project/tables.py

@@ -124,13 +124,27 @@ def headerData(self, section, orientation, role=QtCore.Qt.ItemDataRole.DisplayRo
             else:
                 header = header.replace("_", " ").title()
             return header
+        elif orientation == QtCore.Qt.Orientation.Vertical and role == QtCore.Qt.ItemDataRole.DisplayRole:
+            return f"{section + 1}"
         return None
 
     def append_item(self):
         """Append an item to the ClassList."""
         self.classlist.append(self.item_type())
         self.endResetModel()
 
+    def insert_item(self, row: int):
+        """Insert an item in the ClassList at given row.
+
+        Parameters
+        ----------
+        row : int
+            The row to insert the item.
+
+        """
+        self.classlist.insert(row, self.item_type())
+        self.endResetModel()
+
     def delete_item(self, row: int):
         """Delete an item in the ClassList.
 
@@ -188,26 +202,46 @@ def __init__(self, field: str, parent):
         self.parent = parent
         self.project_widget = parent.parent
         self.table = QtWidgets.QTableView(parent)
-
+        self.table.setSelectionMode(self.table.SelectionMode.SingleSelection)
+        self.table.setSelectionBehavior(self.table.SelectionBehavior.SelectItems)
+        self.table.verticalHeader().sectionClicked.connect(self.toggle_row_selection)
         self.table.horizontalHeader().setCascadingSectionResizes(True)
+        self.table.horizontalHeader().setHighlightSections(False)
         self.table.setMinimumHeight(100)
 
         layout = QtWidgets.QVBoxLayout()
         layout.setContentsMargins(0, 0, 0, 0)
-        topbar = QtWidgets.QHBoxLayout()
-        topbar.addWidget(QtWidgets.QLabel(header, objectName="ProjectFieldWidgetLabel"))
+        top_bar = QtWidgets.QHBoxLayout()
+        top_bar.addWidget(QtWidgets.QLabel(header, objectName="ProjectFieldWidgetLabel"))
         self.add_button = QtWidgets.QPushButton(
             f"Add new {header[:-1] if header[-1] == 's' else header}", objectName="ProjectFieldWidgetButton"
         )
         self.add_button.setHidden(True)
         self.add_button.pressed.connect(self.append_item)
-        topbar.addStretch(1)
-        topbar.addWidget(self.add_button)
+        top_bar.addStretch(1)
+        top_bar.addWidget(self.add_button)
 
-        layout.addLayout(topbar)
+        layout.addLayout(top_bar)
         layout.addWidget(self.table)
         self.setLayout(layout)
 
+    def toggle_row_selection(self, index):
+        """Toggle selection of a given row in the table.
+
+        Parameters
+        ----------
+        index : int
+            The row to be deleted.
+
+        """
+        selection = self.table.selectionModel()
+        if selection.isRowSelected(index):
+            selection.clear()
+        else:
+            selection.clear()
+            for i in range(self.model.columnCount()):
+                selection.select(self.model.index(index, i), QtCore.QItemSelectionModel.SelectionFlag.Select)
+
     def resizeEvent(self, event):
         self.resize_columns()
         super().resizeEvent(event)
@@ -271,11 +305,19 @@ def append_item(self):
         """Append an item to the model if the model exists."""
         self.model.rowCount()
         if self.model is not None:
-            self.model.append_item()
+            selection = self.table.selectionModel().selectedRows()
+            if selection:
+                cur_row = selection[-1].row() + 1
+                self.model.insert_item(selection[-1].row() + 1)
+            else:
+                self.model.append_item()
+                cur_row = self.model.rowCount() - 1
+                self.table.scrollToBottom()
+            cur_col = self.model.headers.index("name") + self.model.col_offset
+            self.table.setCurrentIndex(self.model.index(cur_row, cur_col))
 
         # call edit again to recreate delete buttons
         self.edit()
-        self.table.scrollToBottom()
 
     def delete_item(self, index):
         """Delete an item at the index if the model exists.
@@ -303,7 +345,14 @@ def edit(self):
         self.resize_columns()
 
     def make_delete_button(self, index):
-        """Make a button that deletes index `index` from the list."""
+        """Make a button that deletes the given row from the list when clicked.
+
+        Parameters
+        ----------
+        index : int
+            The row to be deleted.
+
+        """
         button = QtWidgets.QPushButton(icon=QtGui.QIcon(path_for("delete-dark.png")))
         button.resize(button.sizeHint().width(), button.sizeHint().width())
         button.pressed.connect(lambda: self.delete_item(index))
@@ -340,6 +389,16 @@ def __init__(self, classlist: ratapi.ClassList, parent: QtWidgets.QWidget):
                 if isinstance(item, ratapi.models.ProtectedParameter):
                     self.protected_indices.append(i)
 
+    def data(self, index, role=QtCore.Qt.ItemDataRole.DisplayRole):
+        if not index.isValid():
+            return QtCore.QVariant()
+
+        if role == QtCore.Qt.ItemDataRole.BackgroundRole:
+            header = self.index_header(index)
+            if self.classlist[index.row()].prior_type != "gaussian" and header in ["mu", "sigma"]:
+                return QtGui.QBrush(self.parent.palette().window().color())
+        return super().data(index, role)
+
     def flags(self, index):
         flags = super().flags(index)
         header = self.index_header(index)
@@ -458,6 +517,13 @@ def append_item(self):
         self.classlist.append(self.item_type(**kwargs))
         self.endResetModel()
 
+    def insert_item(self, row: int):
+        kwargs = {"thickness": "", "SLD": "", "roughness": ""}
+        if self.absorption:
+            kwargs["SLD_imaginary"] = ""
+        self.classlist.insert(row, self.item_type(**kwargs))
+        self.endResetModel()
+
     def set_absorption(self, absorption: bool):
         """Set whether the project is using absorption or not.
 

tests/widgets/project/test_models.py

@@ -138,18 +138,32 @@ def test_model_set_data(table_model):
 def test_append(table_model):
     """Test that append_item successfully adds an item of the relevant type."""
     model = table_model
-
     model.append_item()
 
     assert len(model.classlist) == 4
     assert model.classlist[-1].name == "Test Model"
     assert model.classlist[-1].value == 15
 
 
+def test_insert(table_model):
+    """Test that insert_item successfully inserts an item of the relevant type."""
+    model = table_model
+    model.insert_item(1)
+
+    assert len(model.classlist) == 4
+    assert model.classlist[1].name == "Test Model"
+    assert model.classlist[1].value == 15
+
+    model.classlist[1].name = "D"
+    model.insert_item(3)
+    assert len(model.classlist) == 5
+    assert model.classlist[3].name == "Test Model"
+    assert model.classlist[3].value == 15
+
+
 def test_delete(table_model):
     """Test that delete_item deletes the item at the desired index."""
     model = table_model
-
     model.delete_item(1)
 
     assert len(model.classlist) == 2