Merge pull request #30 from AllenNeuralDynamics/bug-fix-harp-device-i…

…nheritance

Make HarpDevice base class inherit from Device base class

docs/conf.py

@@ -11,9 +11,10 @@
 import sys
 
 sys.path.insert(0, os.path.abspath("../src/DataSchemas"))
-from aind_behavior_services import __version__
 import glob
 
+from aind_behavior_services import __version__
+
 SOURCE_ROOT = "https://github.com/AllenNeuralDynamics/Aind.Behavior.Services/tree/main/src/DataSchemas/"
 
 

examples/olfactometer.py

@@ -28,7 +28,7 @@
     ),
     olf.OlfactometerChannel.Channel3: olf.OlfactometerChannelConfig(
         channel_index=olf.OlfactometerChannel.Channel3, channel_type=olf.OlfactometerChannelType.CARRIER, odorant="Air"
-    )
+    ),
 }
 
 calibration = olf.OlfactometerCalibration(

pyproject.toml

@@ -84,7 +84,6 @@ exclude_lines = ["if __name__ == .__main__.:", "pragma: no cover"]
 fail_under = 100
 
 [tool.isort]
-line_length = 120
 profile = "black"
 
 [tool.interrogate]

scripts/regenerate.py

@@ -1,14 +1,17 @@
-from pathlib import Path
 import inspect
+from pathlib import Path
 
-from aind_behavior_services.session import AindBehaviorSessionModel
+from aind_behavior_services import db_utils
+from aind_behavior_services.calibration import aind_manipulator as m
 from aind_behavior_services.calibration import load_cells as lc
 from aind_behavior_services.calibration import olfactometer as olf
 from aind_behavior_services.calibration import water_valve as wv
-from aind_behavior_services.calibration import aind_manipulator as m
-from aind_behavior_services import db_utils
-from aind_behavior_services.utils import convert_pydantic_to_bonsai, pascal_to_snake_case, snake_to_pascal_case
-
+from aind_behavior_services.session import AindBehaviorSessionModel
+from aind_behavior_services.utils import (
+    convert_pydantic_to_bonsai,
+    pascal_to_snake_case,
+    snake_to_pascal_case,
+)
 
 SCHEMA_ROOT = Path("./src/DataSchemas/schemas")
 EXTENSIONS_ROOT = Path("./src/Extensions/")

src/DataSchemas/aind_behavior_services/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.7.3"
+__version__ = "0.7.4"
 
 from .rig import AindBehaviorRigModel
 from .session import AindBehaviorSessionModel

src/DataSchemas/aind_behavior_services/aind_data_schema_mapper/mapper.py

@@ -3,8 +3,18 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Type, TypeVar, Union
 
-from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
-from aind_data_schema.core.session import Modality, Session, Software, StimulusEpoch, Stream
+from aind_behavior_services import (
+    AindBehaviorRigModel,
+    AindBehaviorSessionModel,
+    AindBehaviorTaskLogicModel,
+)
+from aind_data_schema.core.session import (
+    Modality,
+    Session,
+    Software,
+    StimulusEpoch,
+    Stream,
+)
 
 TSession = TypeVar("TSession", bound=AindBehaviorSessionModel)
 TRig = TypeVar("TRig", bound=AindBehaviorRigModel)

src/DataSchemas/aind_behavior_services/calibration/load_cells.py

@@ -62,6 +62,7 @@ class CalibrationParameters(TaskParameters):
 
 class CalibrationLogic(AindBehaviorTaskLogicModel):
     """Load cells operation control model that is used to run a calibration data acquisition workflow"""
+
     name: str = Field(default="LoadCellsCalibrationLogic", title="Task name")
     version: Literal[TASK_LOGIC_VERSION] = TASK_LOGIC_VERSION
     task_parameters: CalibrationParameters = Field(..., title="Task parameters", validate_default=True)

src/DataSchemas/aind_behavior_services/calibration/olfactometer.py

@@ -2,7 +2,12 @@
 from typing import Dict, Literal, Optional
 
 from aind_behavior_services.calibration import Calibration
-from aind_behavior_services.rig import AindBehaviorRigModel, HarpAnalogInput, HarpClockGenerator, HarpOlfactometer
+from aind_behavior_services.rig import (
+    AindBehaviorRigModel,
+    HarpAnalogInput,
+    HarpClockGenerator,
+    HarpOlfactometer,
+)
 from aind_behavior_services.task_logic import AindBehaviorTaskLogicModel, TaskParameters
 from pydantic import BaseModel, Field
 
@@ -68,6 +73,7 @@ class CalibrationParameters(TaskParameters):
 
 class CalibrationLogic(AindBehaviorTaskLogicModel):
     """Olfactometer operation control model that is used to run a calibration data acquisition workflow"""
+
     name: str = Field(default="OlfactometerCalibrationLogic", title="Task name")
     version: Literal[TASK_LOGIC_VERSION] = TASK_LOGIC_VERSION
     task_parameters: CalibrationParameters = Field(..., title="Task parameters", validate_default=True)

src/DataSchemas/aind_behavior_services/calibration/water_valve.py

@@ -131,6 +131,7 @@ class CalibrationParameters(TaskParameters):
 
 class CalibrationLogic(AindBehaviorTaskLogicModel):
     """Olfactometer operation control model that is used to run a calibration data acquisition workflow"""
+
     name: str = Field(default="WaterValveCalibrationLogic", title="Task name")
     version: Literal[TASK_LOGIC_VERSION] = TASK_LOGIC_VERSION
     task_parameters: CalibrationParameters = Field(..., title="Task parameters", validate_default=True)

src/DataSchemas/aind_behavior_services/launcher.py

@@ -5,7 +5,11 @@
 from typing import Generic, List, Optional, Tuple, Type, TypeVar, Union
 
 import git
-from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
+from aind_behavior_services import (
+    AindBehaviorRigModel,
+    AindBehaviorSessionModel,
+    AindBehaviorTaskLogicModel,
+)
 from aind_behavior_services.db_utils import SubjectDataBase, SubjectEntry
 from aind_behavior_services.utils import open_bonsai_process
 from pydantic import ValidationError
@@ -274,7 +278,7 @@ def prompt_session_input(self, folder: Optional[str] = None) -> TSession:
         notes = self._get_notes()
 
         return self.session_schema(
-            experiment="", # Will be set later
+            experiment="",  # Will be set later
             root_path=self.data_dir,
             remote_path=self.remote_data_dir,
             subject=subject,

src/DataSchemas/aind_behavior_services/rig/__init__.py

@@ -80,69 +80,69 @@ class HarpDeviceType(str, Enum):
     GENERIC = "generic"
 
 
-class HarpDeviceBase(BaseModel):
+class HarpDeviceGeneric(Device):
     who_am_i: Optional[int] = Field(default=None, le=9999, ge=0, description="Device WhoAmI")
-    device_type: HarpDeviceType = Field(default=HarpDeviceType.GENERIC, description="Device type")
+    device_type: Literal[HarpDeviceType.GENERIC] = HarpDeviceType.GENERIC
     serial_number: Optional[str] = Field(default=None, description="Device serial number")
     port_name: str = Field(..., description="Device port name")
 
 
-class HarpBehavior(HarpDeviceBase):
+class HarpBehavior(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.BEHAVIOR] = HarpDeviceType.BEHAVIOR
     who_am_i: Literal[1216] = 1216
 
 
-class HarpSoundCard(HarpDeviceBase):
+class HarpSoundCard(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.SOUNDCARD] = HarpDeviceType.SOUNDCARD
     who_am_i: Literal[1280] = 1280
 
 
-class HarpLoadCells(HarpDeviceBase):
+class HarpLoadCells(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.LOADCELLS] = HarpDeviceType.LOADCELLS
     who_am_i: Literal[1232] = 1232
 
 
-class HarpOlfactometer(HarpDeviceBase):
+class HarpOlfactometer(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.OLFACTOMETER] = HarpDeviceType.OLFACTOMETER
     who_am_i: Literal[1140] = 1140
 
 
-class HarpClockGenerator(HarpDeviceBase):
+class HarpClockGenerator(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.CLOCKGENERATOR] = HarpDeviceType.CLOCKGENERATOR
     who_am_i: Literal[1158] = 1158
 
 
-class HarpClockSynchronizer(HarpDeviceBase):
+class HarpClockSynchronizer(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.CLOCKSYNCHRONIZER] = HarpDeviceType.CLOCKSYNCHRONIZER
     who_am_i: Literal[1152] = 1152
 
 
-class HarpAnalogInput(HarpDeviceBase):
+class HarpAnalogInput(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.ANALOGINPUT] = HarpDeviceType.ANALOGINPUT
     who_am_i: Literal[1236] = 1236
 
 
-class HarpLickometer(HarpDeviceBase):
+class HarpLickometer(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.LICKOMETER] = HarpDeviceType.LICKOMETER
     who_am_i: Literal[1400] = 1400
 
 
-class HarpSniffDetector(HarpDeviceBase):
+class HarpSniffDetector(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.SNIFFDETECTOR] = HarpDeviceType.SNIFFDETECTOR
     who_am_i: Literal[1401] = 1401
 
 
-class HarpTreadmill(HarpDeviceBase):
+class HarpTreadmill(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.TREADMILL] = HarpDeviceType.TREADMILL
     who_am_i: Literal[1402] = 1402
 
 
-class HarpCuttlefish(HarpDeviceBase):
+class HarpCuttlefish(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.CUTTLEFISH] = HarpDeviceType.CUTTLEFISH
     who_am_i: Literal[1403] = 1403
 
 
-class HarpStepperDriver(HarpDeviceBase):
+class HarpStepperDriver(HarpDeviceGeneric):
     device_type: Literal[HarpDeviceType.STEPPERDRIVER] = HarpDeviceType.STEPPERDRIVER
     who_am_i: Literal[1130] = 1130
 

src/DataSchemas/aind_behavior_services/utils.py

@@ -9,7 +9,12 @@
 from typing import Dict, List, Optional, TypeVar
 
 from pydantic import BaseModel, PydanticInvalidForJsonSchema
-from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue, _deduplicate_schemas
+from pydantic.json_schema import (
+    GenerateJsonSchema,
+    JsonSchemaMode,
+    JsonSchemaValue,
+    _deduplicate_schemas,
+)
 from pydantic_core import PydanticOmit, core_schema, to_jsonable_python
 
 

tests/test_launcher.py

@@ -1,7 +1,11 @@
 import unittest
 from typing import Literal
 
-from aind_behavior_services import AindBehaviorRigModel, AindBehaviorSessionModel, AindBehaviorTaskLogicModel
+from aind_behavior_services import (
+    AindBehaviorRigModel,
+    AindBehaviorSessionModel,
+    AindBehaviorTaskLogicModel,
+)
 from aind_behavior_services.launcher import Launcher, LauncherCli
 from pydantic import Field
 

tests/test_schema_version_coercion.py

@@ -4,7 +4,7 @@
 
 from aind_behavior_services import AindBehaviorTaskLogicModel
 from aind_behavior_services.task_logic import TaskParameters
-from pydantic import ValidationError, Field
+from pydantic import Field, ValidationError
 
 version_pre: str = "0.0.1"
 version_post: str = "0.0.2"
@@ -15,6 +15,7 @@ class AindBehaviorRigModelPre(AindBehaviorTaskLogicModel):
     name: str = Field(default="Pre")
     task_parameters: TaskParameters = Field(default=TaskParameters(), validate_default=True)
 
+
 class AindBehaviorRigModelPost(AindBehaviorTaskLogicModel):
     version: Literal[version_post] = version_post
     name: str = Field(default="Post")
@@ -34,9 +35,7 @@ def test_version_update_coercion(self):
             except ValidationError as e:
                 self.fail(f"Validation failed with error: {e}")
 
-            self.assertEqual(
-                pre_updated.version, post_instance.version, "Schema version was not coerced correctly."
-            )
+            self.assertEqual(pre_updated.version, post_instance.version, "Schema version was not coerced correctly.")
 
 
 if __name__ == "__main__":