From b5f885f0a40a4b1a24ab5bd384382cf2aa377a45 Mon Sep 17 00:00:00 2001
From: Miles Wells <k1o0@3tk.co>
Date: Tue, 5 Mar 2024 16:40:09 -0500
Subject: [PATCH] separate task dependancies

---
 iblrig_custom_tasks/_sp_passiveVideo/task.py  | 11 ++++++--
 .../samuel_cuedBiasedChoiceWorld/task.py      |  9 +++----
 projects/neuromodulators.py                   | 26 +++++++++++--------
 pyproject.toml                                |  4 ++-
 4 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/iblrig_custom_tasks/_sp_passiveVideo/task.py b/iblrig_custom_tasks/_sp_passiveVideo/task.py
index 98c1ead..c7e4eec 100644
--- a/iblrig_custom_tasks/_sp_passiveVideo/task.py
+++ b/iblrig_custom_tasks/_sp_passiveVideo/task.py
@@ -7,8 +7,15 @@
 from pathlib import Path
 from collections import defaultdict
 import logging
-
-import vlc
+import warnings
+
+try:
+    import vlc
+except ModuleNotFoundError:
+    warnings.warn(
+        'Please install extra dependencies for _sp_passiveVideo: '
+        'pip install "project_extraction[_sp_passiveVideo] @ '
+        'git+https://github.com/int-brain-lab/project_extraction.git"', RuntimeWarning)
 import pandas as pd
 from pybpodapi.protocol import Bpod
 
diff --git a/iblrig_custom_tasks/samuel_cuedBiasedChoiceWorld/task.py b/iblrig_custom_tasks/samuel_cuedBiasedChoiceWorld/task.py
index 0d18383..e7566a2 100644
--- a/iblrig_custom_tasks/samuel_cuedBiasedChoiceWorld/task.py
+++ b/iblrig_custom_tasks/samuel_cuedBiasedChoiceWorld/task.py
@@ -3,7 +3,6 @@
 from pybpodapi.protocol import StateMachine
 
 import iblrig.misc
-from iblrig.path_helper import load_pydantic_yaml, HardwareSettings
 from iblrig.base_choice_world import BiasedChoiceWorldSession
 from iblrig.hardware import SOFTCODE
 from iblutil.util import setup_logger
@@ -19,15 +18,13 @@ class Session(BiasedChoiceWorldSession):
     protocol_name = 'samuel_cuedBiasedChoiceWorld'
 
     def __init__(self, *args, delay_secs=0, **kwargs):  #SP _init_ should be the same as biasedChoiceWorld, so should it be specified?
+        super().__init__(**kwargs)
+
         # loads in the settings in order to determine the main sync and thus the pipeline extractor tasks
-        hardware_settings = load_pydantic_yaml(HardwareSettings, kwargs.get('file_hardware_settings'))
-        hardware_settings.update(kwargs.get('hardware_settings', {}))
-        is_main_sync = hardware_settings.get('MAIN_SYNC', False)
+        is_main_sync = self.hardware_settings.get('MAIN_SYNC', False)
         trials_task = 'CuedBiasedTrials' if is_main_sync else 'CuedBiasedTrialsTimeline'
         self.extractor_tasks = ['TrialRegisterRaw', trials_task, 'TrainingStatus']
 
-        super().__init__(**kwargs)
-
         self.task_params["SESSION_DELAY_START"] = delay_secs
         # init behaviour data
         self.movement_left = self.device_rotary_encoder.THRESHOLD_EVENTS[
diff --git a/projects/neuromodulators.py b/projects/neuromodulators.py
index 8585877..3df5604 100644
--- a/projects/neuromodulators.py
+++ b/projects/neuromodulators.py
@@ -1,20 +1,24 @@
 import numpy as np
 from ibllib.io.extractors.biased_trials import BiasedTrials
+from ibllib.io.extractors.base import BaseBpodTrialsExtractor, run_extractor_classes
 
 
-class TrialsTableNeuromodulator(BiasedTrials):
-    def __init__(self, *args, **kwargs):
-        """Include custom trial variables in `var_names` property.
+class TrialsTableNeuromodulator(BaseBpodTrialsExtractor):
+    """Extract neuromodulator task events.
 
-        NB: In order to save to file and register these variables, we would need a custom ibllib.pipes.tasks.Task with
-        the correct output file names.
-        """
-        super().__init__(*args, **kwargs)
-        self.save_names = BiasedTrials.save_names + (None, None)
-        self.var_names = BiasedTrials.var_names + ('omit_feedback', 'exit_state')
+    Include a couple custom trial variables in `var_names` property.
+
+    NB: In order to save to file and register these variables, we would need a custom ibllib.pipes.tasks.Task with
+    the correct output file names.
+    """
+    save_names = BiasedTrials.save_names + (None, None)
+    var_names = BiasedTrials.var_names + ('omit_feedback', 'exit_state')
 
     def _extract(self, *args, **kwargs):
-        out = super(TrialsTableNeuromodulator, self)._extract(*args, **kwargs)
+        # Extract common biased choice world datasets
+        out, _ = run_extractor_classes(
+            [BiasedTrials], session_path=self.session_path, bpod_trials=self.bpod_trials,
+            settings=self.settings, save=False, task_collection=self.task_collection)
         out[0]['omit_feedback'] = np.array([t['omit_feedback'] for t in self.bpod_trials])
         out[0]['exit_state'] = np.array([t['behavior_data']['States timestamps']['exit_state'][0][0] for t in self.bpod_trials])
-        return out
+        return {k: out[k] for k in self.var_names}  # Ensures all datasets present and ordered
diff --git a/pyproject.toml b/pyproject.toml
index f461ff4..3509381 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,9 @@ dynamic         = [ "readme" ]
 keywords        = [ "IBL", "neuro-science" ]
 requires-python = "~=3.10"
 license         = { file = "LICENSE" }
-dependencies    = [ "python-vlc" ]
+
+[project.optional-dependencies]
+_sp_passiveVideo    = [ "python-vlc" ]
 
 [tool.setuptools]
 include-package-data = true