Merge pull request #49 from AllenNeuralDynamics:feat-add-types

Add data types

scripts/regenerate.py

@@ -5,6 +5,7 @@
 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.data_types import DataTypes
 from aind_behavior_services.session import AindBehaviorSessionModel
 from aind_behavior_services.utils import (
     convert_pydantic_to_bonsai,
@@ -48,6 +49,11 @@ def main():
         namespace=f"{NAMESPACE_PREFIX}.AindBehaviorSession",
     )
 
+    convert_pydantic_to_bonsai(
+        {"aind_behavior_data_types": DataTypes},
+        schema_path=SCHEMA_ROOT, skip_sgen=True
+    )
+
     convert_pydantic_to_bonsai(
         {"aind_behavior_subject_database": db_utils.SubjectDataBase},
         schema_path=SCHEMA_ROOT,

src/DataSchemas/aind_behavior_services/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.7.11"
+__version__ = "0.7.12"
 
 from .rig import AindBehaviorRigModel  # noqa: F401
 from .session import AindBehaviorSessionModel  # noqa: F401

src/DataSchemas/aind_behavior_services/data_types.py

@@ -0,0 +1,51 @@
+from enum import StrEnum
+from typing import Any, Optional, Literal
+from aind_behavior_services.base import SchemaVersionedModel
+from pydantic import BaseModel, Field
+
+__version__ = "0.1.0"
+
+class DataType(StrEnum):
+    STRING = "string"
+    NUMBER = "number"
+    BOOLEAL = "boolean"
+    OBJECT = "object"
+    ARRAY = "array"
+    NULL = "null"
+
+
+class TimestampSource(StrEnum):
+    NULL = "null"
+    HARP = "harp"
+    RENDER = "render"
+
+
+class SoftwareEvent(BaseModel):
+    """
+    A software event is a generic event that can be used to track any event that occurs in the software.
+    """
+
+    name: str = Field(..., description="The name of the event")
+    timestamp: Optional[float] = Field(..., description="The timestamp of the event")
+    timestamp_source: TimestampSource = Field(default=TimestampSource.NULL, description="The source of the timestamp")
+    frame_index: Optional[int] = Field(default=None, ge=0, description="The frame index of the event")
+    frame_timestamp: Optional[float] = Field(default=None, description="The timestamp of the frame")
+    data: Optional[Any] = Field(..., description="The data of the event")
+    data_type: DataType = Field(default=DataType.NULL, alias="dataType", description="The data type of the event")
+
+
+class RenderSynchState(BaseModel):
+    sync_quad_value: Optional[float] = Field(default=None, ge=0, le=1, description="The synchronization quad value")
+    frame_index: Optional[int] = Field(default=None, ge=0, description="The frame index of the event")
+    frame_timestamp: Optional[float] = Field(default=None, ge=0, description="The timestamp of the frame")
+
+
+class DataTypes(SchemaVersionedModel):
+    version: Literal[__version__] = __version__
+    software_event: SoftwareEvent
+    render_synch_state: RenderSynchState
+
+    class Config:
+        json_schema_extra = {
+            "x-abstract": "True",
+        }

src/DataSchemas/aind_behavior_services/utils.py

@@ -7,14 +7,15 @@
 from pathlib import Path
 from string import capwords
 from subprocess import CompletedProcess, run
-from typing import Any, Dict, List, Optional, TypeVar
+from typing import Any, Dict, List, NewType, Optional, Type, Union
 
 from pydantic import BaseModel, PydanticInvalidForJsonSchema
 from pydantic.json_schema import (
     GenerateJsonSchema,
     JsonSchemaMode,
     JsonSchemaValue,
     _deduplicate_schemas,
+    models_json_schema,
 )
 from pydantic_core import PydanticOmit, core_schema, to_jsonable_python
 
@@ -167,17 +168,22 @@ def union_schema(self, schema: core_schema.UnionSchema) -> JsonSchemaValue:
             return self.get_flattened_anyof(generated)
 
 
-Model = TypeVar("Model", bound=BaseModel)
+ModelInputTypeSignature = Union[List[Type[BaseModel]] | Type[BaseModel]]
 
 
 def export_schema(
-    model: Model,
-    schema_generator: GenerateJsonSchema = CustomGenerateJsonSchema,
+    model: ModelInputTypeSignature,
+    schema_generator: Type[GenerateJsonSchema] = CustomGenerateJsonSchema,
     mode: JsonSchemaMode = "serialization",
     def_keyword: str = "definitions",
+    models_title: Optional[str] = None,
 ):
     """Export the schema of a model to a json file"""
-    _model = model.model_json_schema(schema_generator=schema_generator, mode=mode)
+    if not isinstance(model, list):
+        _model = model.model_json_schema(schema_generator=schema_generator, mode=mode)
+    else:
+        models = [(m, mode) for m in model]
+        _, _model = models_json_schema(models, schema_generator=schema_generator, title=models_title)
     json_model = json.dumps(_model, indent=2)
     json_model = json_model.replace("$defs", def_keyword)
     return json_model
@@ -233,27 +239,30 @@ def bonsai_sgen(
 
 
 def convert_pydantic_to_bonsai(
-    models: Dict[str, BaseModel],
+    models: Dict[str, ModelInputTypeSignature],
     namespace: str = "DataSchema",
     schema_path: PathLike = Path("./src/DataSchemas/"),
     output_path: PathLike = Path("./src/Extensions/"),
     serializer: Optional[List[BonsaiSgenSerializers]] = None,
+    skip_sgen: bool = False,
+    export_schema_kwargs: Dict[str, Any] = {},
 ) -> None:
 
-    def _write_json(schema_path: PathLike, output_model_name: str, model: BaseModel) -> None:
+    def _write_json(schema_path: PathLike, output_model_name: str, model: ModelInputTypeSignature) -> None:
         with open(os.path.join(schema_path, f"{output_model_name}.json"), "w", encoding="utf-8") as f:
-            json_model = export_schema(model)
+            json_model = export_schema(model, **export_schema_kwargs)
             f.write(json_model)
 
     for output_model_name, model in models.items():
         _write_json(schema_path, output_model_name, model)
-        cmd_return = bonsai_sgen(
-            schema_path=Path(os.path.join(schema_path, f"{output_model_name}.json")),
-            output_path=Path(os.path.join(output_path, f"{snake_to_pascal_case(output_model_name)}.cs")),
-            namespace=namespace,
-            serializer=serializer,
-        )
-        print(cmd_return.stdout)
+        if not skip_sgen:
+            cmd_return = bonsai_sgen(
+                schema_path=Path(os.path.join(schema_path, f"{output_model_name}.json")),
+                output_path=Path(os.path.join(output_path, f"{snake_to_pascal_case(output_model_name)}.cs")),
+                namespace=namespace,
+                serializer=serializer,
+            )
+            print(cmd_return.stdout)
 
 
 def snake_to_pascal_case(s: str) -> str:

src/DataSchemas/schemas/aind_behavior_data_types.json

@@ -0,0 +1,179 @@
+{
+  "definitions": {
+    "DataType": {
+      "enum": [
+        "string",
+        "number",
+        "boolean",
+        "object",
+        "array",
+        "null"
+      ],
+      "title": "DataType",
+      "type": "string"
+    },
+    "RenderSynchState": {
+      "properties": {
+        "sync_quad_value": {
+          "default": null,
+          "description": "The synchronization quad value",
+          "oneOf": [
+            {
+              "maximum": 1.0,
+              "minimum": 0.0,
+              "type": "number"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Sync Quad Value"
+        },
+        "frame_index": {
+          "default": null,
+          "description": "The frame index of the event",
+          "oneOf": [
+            {
+              "minimum": 0,
+              "type": "integer"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Frame Index"
+        },
+        "frame_timestamp": {
+          "default": null,
+          "description": "The timestamp of the frame",
+          "oneOf": [
+            {
+              "minimum": 0.0,
+              "type": "number"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Frame Timestamp"
+        }
+      },
+      "title": "RenderSynchState",
+      "type": "object"
+    },
+    "SoftwareEvent": {
+      "description": "A software event is a generic event that can be used to track any event that occurs in the software.",
+      "properties": {
+        "name": {
+          "description": "The name of the event",
+          "title": "Name",
+          "type": "string"
+        },
+        "timestamp": {
+          "description": "The timestamp of the event",
+          "oneOf": [
+            {
+              "type": "number"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Timestamp"
+        },
+        "timestamp_source": {
+          "allOf": [
+            {
+              "$ref": "#/definitions/TimestampSource"
+            }
+          ],
+          "default": "null",
+          "description": "The source of the timestamp"
+        },
+        "frame_index": {
+          "default": null,
+          "description": "The frame index of the event",
+          "oneOf": [
+            {
+              "minimum": 0,
+              "type": "integer"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Frame Index"
+        },
+        "frame_timestamp": {
+          "default": null,
+          "description": "The timestamp of the frame",
+          "oneOf": [
+            {
+              "type": "number"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Frame Timestamp"
+        },
+        "data": {
+          "description": "The data of the event",
+          "oneOf": [
+            {},
+            {
+              "type": "null"
+            }
+          ],
+          "title": "Data"
+        },
+        "dataType": {
+          "allOf": [
+            {
+              "$ref": "#/definitions/DataType"
+            }
+          ],
+          "default": "null",
+          "description": "The data type of the event"
+        }
+      },
+      "required": [
+        "name",
+        "timestamp",
+        "data"
+      ],
+      "title": "SoftwareEvent",
+      "type": "object"
+    },
+    "TimestampSource": {
+      "enum": [
+        "null",
+        "harp",
+        "render"
+      ],
+      "title": "TimestampSource",
+      "type": "string"
+    }
+  },
+  "properties": {
+    "version": {
+      "const": "0.1.0",
+      "default": "0.1.0",
+      "title": "Version",
+      "type": "string"
+    },
+    "software_event": {
+      "$ref": "#/definitions/SoftwareEvent"
+    },
+    "render_synch_state": {
+      "$ref": "#/definitions/RenderSynchState"
+    }
+  },
+  "required": [
+    "software_event",
+    "render_synch_state"
+  ],
+  "title": "DataTypes",
+  "type": "object",
+  "x-abstract": "True"
+}