"""Step configuration module."""
from typing import Any
from typing import ClassVar
from typing import Literal
from typing import cast
from pydantic import Field
from pydantic import model_validator
from paidiverpy.models.colour_params import COLOUR_LAYER_METHODS
from paidiverpy.models.colour_params import ColourParamsUnion
from paidiverpy.models.convert_params import CONVERT_LAYER_METHODS
from paidiverpy.models.convert_params import ConvertParamsUnion
from paidiverpy.models.custom_params import CustomParams
from paidiverpy.models.position_params import POSITION_LAYER_METHODS
from paidiverpy.models.position_params import PositionParamsUnion
from paidiverpy.models.sampling_params import SAMPLING_LAYER_METHODS
from paidiverpy.models.sampling_params import SamplingParamsUnion
from paidiverpy.utils.base_model import BaseModel
steps_params_mapping = {
"colour": COLOUR_LAYER_METHODS,
"convert": CONVERT_LAYER_METHODS,
"position": POSITION_LAYER_METHODS,
"sampling": SAMPLING_LAYER_METHODS,
}
PositionModeLiteral = cast(type, Literal.__getitem__(tuple(POSITION_LAYER_METHODS.keys())))
ColourModeLiteral = cast(type, Literal.__getitem__(tuple(COLOUR_LAYER_METHODS.keys())))
ConvertModeLiteral = cast(type, Literal.__getitem__(tuple(CONVERT_LAYER_METHODS.keys())))
SamplingModeLiteral = cast(type, Literal.__getitem__(tuple(SAMPLING_LAYER_METHODS.keys())))
[docs]
class StepConfig(BaseModel):
"""Step configuration model."""
name: str | None = Field(None, description="Name of the step")
step_name: str | None = Field(None, description="Step name")
test: bool = Field(False, description="Test mode")
params: Any = Field(default_factory=dict, description="Parameters for the step")
mode: PositionModeLiteral | ColourModeLiteral | ConvertModeLiteral | SamplingModeLiteral = Field(
description="Mode for the position step",
)
model_config: ClassVar[dict] = {
"frozen": False,
}
@model_validator(mode="before")
@classmethod
[docs]
def resolve_params_schema(cls, values: dict[str, Any]) -> dict[str, Any]:
"""Resolve the parameters schema based on the step name and mode.
Args:
values (dict): The values to validate.
Returns:
dict: The validated values.
"""
if isinstance(values, StepConfig):
return values
step_name = values.get("step_name")
params: dict[str, Any] = values.get("params", {})
mode = values.get("mode")
if step_name == "custom":
values["params"] = CustomParams(**params)
else:
if step_name not in steps_params_mapping:
msg = f"Unknown step_name: '{step_name}'"
raise ValueError(msg)
if mode is None:
msg = "Missing 'mode' for non-custom step"
raise ValueError(msg)
method_dict = steps_params_mapping[step_name]
if mode not in method_dict:
msg = f"Mode '{mode}' not valid for step '{step_name}'"
raise ValueError(msg)
param_class = method_dict[mode]["params"]
values["params"] = param_class(**params)
return values
[docs]
def update(self, **updates: dict[str, Any]) -> "StepConfig":
"""Update the model in-place with new values."""
for key, value in updates.items():
setattr(self, key, value)
validated = self.__class__.model_validate(self.model_dump())
for key, val in validated.model_dump().items():
setattr(self, key, val)
return self
[docs]
class PositionConfig(StepConfig):
"""Position configuration model."""
name: str | None = Field("position", description="Name of the step")
mode: PositionModeLiteral = Field(
description="Mode for the position step",
)
test: bool = Field(False, description="Test mode")
params: PositionParamsUnion | None = Field(default=None, description="Position parameters")
# processing_type: Literal["image", "dataset"] = Field(
# "image",
# description=("If the images are processed individually (image option) or as a dataset (dataset option)"),
# )
[docs]
class ColourConfig(StepConfig):
"""Colour configuration model."""
name: str | None = Field("colour", description="Name of the step")
mode: ColourModeLiteral = Field(
description="Mode for the colour step",
)
test: bool = Field(False, description="Test mode")
params: ColourParamsUnion | None = Field(default=None, description="Colour parameters")
# processing_type: Literal["image", "dataset"] = Field(
# "image",
# description=("If the images are processed individually (image option) or as a dataset (dataset option)"),
# )
[docs]
class ConvertConfig(StepConfig):
"""Convert configuration model."""
name: str | None = Field("convert", description="Name of the step")
mode: ConvertModeLiteral = Field(description="Mode for the convert step")
test: bool = Field(False, description="Test mode")
params: ConvertParamsUnion | None = Field(default=None, description="Convert parameters")
# processing_type: Literal["image", "dataset"] = Field(
# "image",
# description=("If the images are processed individually (image option) or as a dataset (dataset option)"),
# )
[docs]
class SamplingConfig(StepConfig):
"""Sampling configuration model."""
name: str | None = Field("sampling", description="Name of the step")
mode: SamplingModeLiteral = Field(description="Mode for the sampling step")
test: bool = Field(False, description="Test mode")
params: SamplingParamsUnion | None = Field(default=None, description="Sampling parameters")
# processing_type: Literal["image", "dataset"] = Field(
# "image",
# description=("If the images are processed individually (image option) or as a dataset (dataset option)"),
# )
[docs]
class CustomConfig(BaseModel):
"""Custom configuration model."""
name: str | None = Field("custom", description="Name of the step")
file_path: str = Field(description="File path for custom step")
class_name: str = Field(description="Class name for custom step")
step_name: str | None = Field(None, description="Step name")
test: bool = Field(False, description="Test mode")
processing_type: Literal["image", "dataset"] = Field(
"image",
description=("If the images are processed individually (image option) or as a dataset (dataset option)"),
)
dependencies: str | None = Field(
None,
description=("Dependencies for the custom step. It should be a string with each dependency separated by commas."),
)
dependencies_path: str | None = Field(None, description="Path to the requirements file for the custom step")
params: CustomParams | None = Field(default=None, description="Custom parameters")
StepConfigUnion = ColourConfig | ConvertConfig | CustomConfig | PositionConfig | SamplingConfig
