Allow to init CellMorphologyProtocol directly (#127)

Author: GianlucaFicarelli

examples/02_morphology.ipynb

@@ -22,7 +22,6 @@
     "    CellMorphology,\n",
     "    CellMorphologyProtocol,\n",
     "    Contribution,\n",
-    "    DigitalReconstructionCellMorphologyProtocol,\n",
     "    MTypeClass,\n",
     "    MTypeClassification,\n",
     "    Organization,\n",
@@ -32,6 +31,7 @@
     "    Subject,\n",
     ")\n",
     "from entitysdk.types import (\n",
+    "    CellMorphologyGenerationType,\n",
     "    CellMorphologyProtocolDesign,\n",
     "    SlicingDirectionType,\n",
     ")"
@@ -195,16 +195,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "morphology_protocol = DigitalReconstructionCellMorphologyProtocol(\n",
+    "morphology_protocol = CellMorphologyProtocol(\n",
+    "  generation_type=CellMorphologyGenerationType.digital_reconstruction,\n",
     "  protocol_document=\"https://example.com/\",\n",
     "  protocol_design=CellMorphologyProtocolDesign.cell_patch,\n",
     "  slicing_thickness=20.0,\n",
     "  slicing_direction=SlicingDirectionType.horizontal,\n",
     ")\n",
+    "rprint(morphology_protocol.__class__.__name__)\n",
     "morphology_protocol = client.register_entity(morphology_protocol)\n",
     "rprint(morphology_protocol)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "ff0d1d01",
+   "metadata": {},
+   "source": [
+    "### Optional, list morphology protocols"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71818153",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "protocols = client.search_entity(entity_type=CellMorphologyProtocol).all()\n",
+    "rprint(protocols)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "a0193055-3448-4d1c-99cc-4d953727ef4e",

src/entitysdk/models/__init__.py

@@ -6,13 +6,7 @@
 from entitysdk.models.brain_region import BrainRegion
 from entitysdk.models.brain_region_hierarchy import BrainRegionHierarchy
 from entitysdk.models.cell_morphology import CellMorphology
-from entitysdk.models.cell_morphology_protocol import (
-    CellMorphologyProtocol,
-    ComputationallySynthesizedCellMorphologyProtocol,
-    DigitalReconstructionCellMorphologyProtocol,
-    ModifiedReconstructionCellMorphologyProtocol,
-    PlaceholderCellMorphologyProtocol,
-)
+from entitysdk.models.cell_morphology_protocol import CellMorphologyProtocol
 from entitysdk.models.circuit import Circuit
 from entitysdk.models.classification import ETypeClassification, MTypeClassification
 from entitysdk.models.contribution import Contribution, Role
@@ -54,11 +48,9 @@
     "CellMorphology",
     "CellMorphologyProtocol",
     "Circuit",
-    "ComputationallySynthesizedCellMorphologyProtocol",
     "Consortium",
     "Contribution",
     "Derivation",
-    "DigitalReconstructionCellMorphologyProtocol",
     "ElectricalCellRecording",
     "ElectricalRecordingStimulus",
     "EMCellMesh",
@@ -73,13 +65,11 @@
     "License",
     "MEModel",
     "MEModelCalibrationResult",
-    "ModifiedReconstructionCellMorphologyProtocol",
     "MTypeClass",
     "MTypeClassification",
     "NeuronBlock",
     "Organization",
     "Person",
-    "PlaceholderCellMorphologyProtocol",
     "Publication",
     "Role",
     "ScientificArtifactPublicationLink",

src/entitysdk/models/cell_morphology_protocol.py

@@ -2,8 +2,9 @@
 
 from typing import Annotated, Any, ClassVar, Literal
 
-from pydantic import BaseModel, Field, HttpUrl, TypeAdapter
+from pydantic import ConfigDict, Field, HttpUrl, TypeAdapter
 
+from entitysdk.models.core import Identifiable
 from entitysdk.models.entity import Entity
 from entitysdk.types import (
     CellMorphologyGenerationType,
@@ -16,9 +17,17 @@
 
 
 class CellMorphologyProtocolBase(Entity):
-    """Cell Morphology Protocol Base, used by all the protocols except placeholder."""
+    """Cell Morphology Protocol Base, used by all the protocols."""
+
+    # forbid extra parameters to prevent providing attributes of other classes by mistake
+    model_config = ConfigDict(extra="forbid")
 
     type: EntityType | None = EntityType.cell_morphology_protocol
+
+
+class CellMorphologyProtocolExtendedBase(CellMorphologyProtocolBase):
+    """Cell Morphology Protocol Extended Base, used by all the protocols except placeholder."""
+
     protocol_document: Annotated[
         HttpUrl | None,
         Field(description="URL link to protocol document or publication."),
@@ -30,7 +39,7 @@ class CellMorphologyProtocolBase(Entity):
 
 
 class DigitalReconstructionCellMorphologyProtocol(
-    CellMorphologyProtocolBase,
+    CellMorphologyProtocolExtendedBase,
 ):
     """Experimental morphology method for capturing cell morphology data."""
 
@@ -58,7 +67,7 @@ class DigitalReconstructionCellMorphologyProtocol(
 
 
 class ModifiedReconstructionCellMorphologyProtocol(
-    CellMorphologyProtocolBase,
+    CellMorphologyProtocolExtendedBase,
 ):
     """Modified Reconstruction Cell Morphology Protocol."""
 
@@ -69,7 +78,7 @@ class ModifiedReconstructionCellMorphologyProtocol(
 
 
 class ComputationallySynthesizedCellMorphologyProtocol(
-    CellMorphologyProtocolBase,
+    CellMorphologyProtocolExtendedBase,
 ):
     """Computationally Synthesized Cell Morphology Protocol."""
 
@@ -80,7 +89,7 @@ class ComputationallySynthesizedCellMorphologyProtocol(
 
 
 class PlaceholderCellMorphologyProtocol(
-    Entity,
+    CellMorphologyProtocolBase,
 ):
     """Placeholder Cell Morphology Protocol."""
 
@@ -98,24 +107,31 @@ class PlaceholderCellMorphologyProtocol(
 ]
 
 
-class CellMorphologyProtocol(BaseModel):
+class CellMorphologyProtocol(Identifiable):
     """Polymorphic wrapper for consistent API, to be used for searching and retrieving.
 
     The correct specific protocols are automatically instantiated.
 
-    For the registration it's possible to use any of the specific classes:
+    For the registration it's possible to use this same class, or any of the specific classes:
 
     - `DigitalReconstructionCellMorphologyProtocol`
     - `ModifiedReconstructionCellMorphologyProtocol`
     - `ComputationallySynthesizedCellMorphologyProtocol`
     - `PlaceholderCellMorphologyProtocol`
-
-    or this polymorphic class `CellMorphologyProtocol`, but in that case the instance should be
-    created with `CellMorphologyProtocol.model_validate()` to instantiate the correct object.
     """
 
     _adapter: ClassVar[TypeAdapter] = TypeAdapter(CellMorphologyProtocolUnion)
 
+    def __new__(cls, *args, **kwargs) -> CellMorphologyProtocolUnion:  # type: ignore[misc]
+        """Construct a CellMorphologyProtocol from keyword arguments."""
+        if args:
+            msg = "Positional args not supported, use keyword args instead."
+            raise TypeError(msg)
+        return cls._adapter.validate_python(kwargs)
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Catch-all to satisfy type checkers."""
+
     @classmethod
     def model_validate(cls, obj: Any, *args, **kwargs) -> CellMorphologyProtocolUnion:  # type: ignore[override]
         """Return the correct instance of CellMorphologyProtocolUnion."""

src/entitysdk/route.py

@@ -11,6 +11,7 @@
     "BrainLocation": "brain-location",
     "BrainRegion": "brain-region",
     "BrainRegionHierarchy": "brain-region-hierarchy",
+    "CellMorphology": "cell-morphology",
     "Circuit": "circuit",
     "Consortium": "consortium",
     "Contribution": "contribution",
@@ -35,7 +36,6 @@
     "Organization": "organization",
     "Person": "person",
     "Publication": "publication",
-    "CellMorphology": "cell-morphology",
     "ScientificArtifactPublicationLink": "scientific-artifact-publication-link",
     "Simulation": "simulation",
     "SimulationCampaign": "simulation-campaign",

tests/unit/models/test_cell_morphology_protocol.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import pytest
+from pydantic import ValidationError
 
 from entitysdk.models.cell_morphology_protocol import (
     CellMorphologyProtocol,
@@ -112,6 +113,32 @@ def test_register(request, client, httpx_mock, auth_token, json_data_fixture, mo
     assert registered.model_dump(mode="json", exclude_unset=True) == expected_json
 
 
+@pytest.mark.parametrize(
+    ("json_data_fixture", "model_fixture"),
+    [
+        ("json_digital_reconstruction", "model_digital_reconstruction"),
+        ("json_modified_reconstruction", "model_modified_reconstruction"),
+        ("json_computationally_synthesized", "model_computationally_synthesized"),
+        ("json_placeholder", "model_placeholder"),
+    ],
+)
+def test_adapter(request, json_data_fixture, model_fixture):
+    json_data = request.getfixturevalue(json_data_fixture)
+    model = request.getfixturevalue(model_fixture)
+
+    new_model = Model(**json_data)
+    assert new_model == model
+
+    new_model = Model.model_validate(json_data)
+    assert new_model == model
+
+    with pytest.raises(TypeError, match="Positional args not supported"):
+        Model("name")
+
+    with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
+        Model(generation_type=json_data["generation_type"], invalid_input="invalid")
+
+
 @pytest.mark.parametrize(
     ("json_data_fixture", "model_fixture"),
     [