add weight dlg and warning when closing mid-trial (#502)

iblrig/base_tasks.py

@@ -53,7 +53,7 @@ class BaseSession(ABC):
 
     def __init__(self, subject=None, task_parameter_file=None, file_hardware_settings=None,
                  hardware_settings=None, file_iblrig_settings=None, iblrig_settings=None,
-                 one=None, interactive=True, projects=None, procedures=None, stub=None,
+                 one=None, interactive=True, projects=None, procedures=None, subject_weight_grams=None, stub=None,
                  append=False, log_level='INFO', wizard=False):
         """
         :param subject: The subject nickname. Required.
@@ -66,6 +66,7 @@ def __init__(self, subject=None, task_parameter_file=None, file_hardware_setting
         :param interactive:
         :param projects: An optional list of Alyx protocols.
         :param procedures: An optional list of Alyx procedures.
+        :param subject_weight_grams: weight of the subject
         :param stub: A full path to an experiment description file containing experiment information.
         :param append: bool, if True, append to the latest existing session of the same subject for the same day
         :param fmake: (DEPRECATED) if True, only create the raw_behavior_data folder.
@@ -131,7 +132,7 @@ def __init__(self, subject=None, task_parameter_file=None, file_hardware_setting
             'SESSION_END_TIME': None,
             'SESSION_NUMBER': 0,
             'SUBJECT_NAME': subject,
-            'SUBJECT_WEIGHT': None,
+            'SUBJECT_WEIGHT': subject_weight_grams,
             'TOTAL_WATER_DELIVERED': 0,
         })
         # Executes mixins init methods
@@ -383,7 +384,7 @@ def run(self):
         # this prevents from incrementing endlessly the session number if the hardware fails to connect
         self.start_hardware()
         self.create_session()
-        if self.interactive:
+        if self.session_info.SUBJECT_WEIGHT is None and self.interactive:
             self.session_info.SUBJECT_WEIGHT = graph.numinput(
                 "Subject weighing (gr)", f"{self.session_info.SUBJECT_NAME} weight (gr):", nullable=False)
 

iblrig/gui/wizard.py

@@ -133,6 +133,8 @@ def __init__(self, *args, **kwargs):
         self.uiPushFlush.clicked.connect(self.flush)
         self.uiPushStart.clicked.connect(self.startstop)
         self.uiPushPause.clicked.connect(self.pause)
+        self.uiListProjects.clicked.connect(self.enable_UI_elements)
+        self.uiListProcedures.clicked.connect(self.enable_UI_elements)
         self.uiPushConnect.clicked.connect(self.alyx_connect)
         self.lineEditSubject.textChanged.connect(self._filter_subjects)
         self.running_task_process = None
@@ -150,6 +152,24 @@ def __init__(self, *args, **kwargs):
         self.show()
         QtCore.QTimer.singleShot(1, self.check_for_update)
 
+    def closeEvent(self, event):
+        if self.running_task_process is None:
+            event.accept()
+        else:
+            msgBox = QtWidgets.QMessageBox(parent=self)
+            msgBox.setWindowTitle("Hold on")
+            msgBox.setText("A task is running - do you really want to quit?")
+            msgBox.setStandardButtons(QtWidgets.QMessageBox.No | QtWidgets.QMessageBox.Yes)
+            msgBox.setIcon(QtWidgets.QMessageBox().Question)
+            match msgBox.exec_():
+                case QtWidgets.QMessageBox.No:
+                    event.ignore()
+                case QtWidgets.QMessageBox.Yes:
+                    self.setEnabled(False)
+                    self.repaint()
+                    self.startstop()
+                    event.accept()
+
     def check_for_update(self):
         update_available, remote_version = check_for_updates()
         if update_available == 1:
@@ -215,6 +235,12 @@ def pause(self):
     def startstop(self):
         match self.uiPushStart.text():
             case 'Start':
+                dlg = QtWidgets.QInputDialog()
+                weight, ok = dlg.getDouble(self, 'Subject Weight', 'Subject Weight (g):', value=0, min=0,
+                                           flags=dlg.windowFlags() & ~QtCore.Qt.WindowContextHelpButtonHint)
+                if not ok or weight == 0:
+                    return
+
                 self.controller2model()
                 task = EmptySession(subject=self.model.subject, append=self.uiCheckAppend.isChecked(), wizard=True)
                 self.model.session_folder = task.paths['SESSION_FOLDER']
@@ -232,9 +258,10 @@ def startstop(self):
                     cmd.extend(['--procedures', *self.model.procedures])
                 if self.model.projects:
                     cmd.extend(['--projects', *self.model.projects])
+                cmd.extend(['--weight', f'{weight}'])
+                cmd.append('--wizard')
                 if self.uiCheckAppend.isChecked():
                     cmd.append('--append')
-                cmd.append('--wizard')
                 if self.running_task_process is None:
                     self.running_task_process = subprocess.Popen(cmd)
                 self.uiPushStart.setText('Stop')
@@ -281,6 +308,7 @@ def flush(self):
             print(traceback.format_exc())
             print("Cannot find bpod - is it connected?")
             self.uiPushFlush.setChecked(False)
+            self.uiPushFlush.setStyleSheet('')
             return
 
         if not self.uiPushFlush.isChecked():
@@ -293,7 +321,9 @@ def enable_UI_elements(self):
         is_running = self.uiPushStart.text() == 'Stop'
 
         self.uiPushStart.setEnabled(
-            not self.uiPushFlush.isChecked())
+            not self.uiPushFlush.isChecked()
+            and len(self.uiListProjects.selectedIndexes()) > 0
+            and len(self.uiListProcedures.selectedIndexes()) > 0)
         self.uiPushPause.setEnabled(is_running)
         self.uiPushFlush.setEnabled(not is_running)
         self.uiCheckAppend.setEnabled(not is_running)

iblrig/misc.py

@@ -46,6 +46,7 @@ def _get_task_argument_parser(parents=None):
                         help="long description of what is occurring, something like 'Ephys recording with acute probe(s)'; "
                              "be sure to use the double quote characters to encapsulate the description and a space to separate "
                              "multiple procedures")
+    parser.add_argument('-w', '--weight', type=float, dest='subject_weight_grams', required=False, default=None)
     parser.add_argument('--no-interactive', dest='interactive', action='store_false', default=True)
     parser.add_argument('--append', dest='append', action='store_true', default=False)
     parser.add_argument('--stub', type=Path, help="Path to _ibl_experiment.description.yaml stub file.")