Feature/partial First Pydantic Models for parameter validation

elephant/schemas/schema_statistics.py

@@ -64,7 +64,7 @@ class KernelOptions(Enum):
     spiketrains: Any = Field(..., description="Input spike train(s)")
     sampling_period: Any = Field(..., gt=0, description="Time stamp resolution of spike times")
     kernel: Union[KernelOptions, Any] = Field(KernelOptions.auto, description="Kernel for convolution")
-    cutoff: Optional[float] = Field(5.0, gt=0, description="cutoff of probability distribution")
+    cutoff: Optional[float] = Field(5.0, ge=0, description="cutoff of probability distribution")
     t_start: Optional[Any] = Field(None, ge=0, description="Start time")
     t_stop: Optional[Any] = Field(None, gt=0, description="Stop time")
     trim: Optional[bool] = Field(False, description="Only return region of convolved signal")
@@ -77,9 +77,11 @@ class KernelOptions(Enum):
     @field_validator("spiketrains")
     @classmethod
     def validate_spiketrains(cls, v, info):
-        if(isinstance(v, list)):
-            return fv.validate_spiketrains(v, info, allowed_types=(list,), allowed_content_types=(neo.SpikeTrain,))
-        return fv.validate_spiketrain(v, info, allowed_types=(neo.SpikeTrain, elephant.trials.Trials))
+        if(isinstance(v, (list, neo.core.spiketrainlist.SpikeTrainList))):
+            return fv.validate_spiketrains(v, info, allowed_types=(list, neo.core.spiketrainlist.SpikeTrainList), allowed_content_types=(neo.SpikeTrain,))
+        if(isinstance(v, neo.SpikeTrain)):
+            return fv.validate_spiketrain(v, info, allowed_types=(neo.SpikeTrain,))
+        return fv.validate_spiketrains_matrix(v, info)
 
     @field_validator("sampling_period")
     @classmethod
@@ -89,7 +91,9 @@ def validate_quantity(cls, v, info):
     @field_validator("kernel")
     @classmethod
     def validate_kernel(cls, v, info):
-        return fv.validate_type(v, info, allowed_types=(cls.KernelOptions, Kernel), allow_none=False)
+        if v == cls.KernelOptions.auto.value:
+            return v
+        return fv.validate_type(v, info, allowed_types=(Kernel), allow_none=False)
 
     @field_validator("t_start", "t_stop")
     @classmethod
@@ -100,7 +104,8 @@ def validate_time(cls, v, info):
     def validate_model(self) -> Self:             
         if(isinstance(self.kernel, Kernel) and self.cutoff < self.kernel.min_cutoff):
             warnings.warn(f"cutoff {self.cutoff} is smaller than the minimum cutoff {self.kernel.min_cutoff} of the kernel", UserWarning)
-        fv.model_validate_spiketrains_same_t_start_stop(self.spiketrains, self.t_start, self.t_stop, warning=True)
+        if isinstance(self.spiketrains, list):
+            fv.model_validate_spiketrains_same_t_start_stop(self.spiketrains, self.t_start, self.t_stop, warning=True)
         return self
 
 class PydanticTimeHistogram(BaseModel):
@@ -185,13 +190,13 @@ class NanPolicyOptions(Enum):
         omit = "omit"
         _raise = "raise"
 
-    a: Any = Field(..., description="Input array")
+    args: Any = Field(..., description="Input array")
     axis: Union[int, None] = Field(0, description="Compute statistic axis")
     nan_policy: NanPolicyOptions = Field(NanPolicyOptions.propagate, description="How handle input NaNs")
     ddof: Optional[int] = Field(0, ge=0, description="Delta Degrees Of Freedom")
     keepdims: Optional[bool] = Field(False, description="leave reduced axes in one-dimensional result")
 
-    @field_validator("a")
+    @field_validator("args")
     @classmethod
     def validate_array(cls, v, info):
         return fv.validate_array(v, info)
@@ -281,7 +286,7 @@ def serialize_quantity(self, v):
     @field_validator("spiketrains")
     @classmethod
     def validate_spiketrains(cls, v, info):
-        return fv.validate_spiketrains(v, info)
+        return fv.validate_spiketrains(v, info, min_length=0)
 
     @field_validator("warn_tolerance")
     @classmethod