Source code for vibe.config_validation.pydantic_models.vibe_config_model

from typing_extensions import Self
from pydantic import BaseModel, Field, model_validator, field_validator, ConfigDict




class VibeConfigModel(BaseModel):
    """
    The Pydantic Model for the VIBE config.toml file. Here we can define
    any configuration variables along with

    - default behaviour
    - regex string pattern matching for paths and light releases

    We can also define additional model validation intop of the type hints
    when defining attributes. In this VibeConfig model we
    use the @model_validator(mode='after') decorator to tell pydantic
    to run the decorated method after the VibeConfig model has been initially
    validated. For example:

    .. code-block:: python

        @model_validator(mode='after')
        def check_modes_to_run(self) -> Self:
            ModesToRunConfig(**self.modes_to_run)
            return self

    Here we are validating the `modes_to_run` attribute by using the after decorator.
    The validation is actually just calling a second Pydantic Model specifically made
    to validate the `modes_to_run` structure. If the ModesToRun model is successful,
    then we return self.

    """

    base_path: str = Field(pattern="^\/(?:[^\/]+\/)*[^\/]+$", description="Base path for all outputs of VIBE.")

    b2monitoring_path: str = Field(
        pattern="^\/(?:[^\/]+\/)*[^\/]+$", description="Mirabelle database path for MirabelleUploadTask."
    )

    analysis_validation_light_release: str = Field(
        default="light-2503-ceres", description="basf2 light release name, must contain 'light'."
    )

    gbasf2_submission_campaign: str = Field(
        pattern=".*v.*", description="Submission campaign suffix for gbasf2 project name."
    )

    use_belle_dataprod_proxy_group: bool = Field(
        default=False, description="If True, use the dataprod group proxy for grid submission."
    )

    luigi_workers: int = Field(gt=0, description="Number of Luigi workers to run.")

    modes_to_run: dict[str, dict] = Field(description="Structure defining the modes which need to be run.")



    @model_validator(mode="after")
    def check_modes_to_run(self) -> Self:
        ModesToRunConfig(**self.modes_to_run)
        return self






class ModesToRunConfig(BaseModel):
    """
    Model for validating the `modes_to_run` variables structure defined in the config.toml.
    Each attribute maps to one of the ModeTypes defined in the mode.py module and all have identical
    value structures:

    .. code-block:: python

        class ModesToRunConfig(BaseModel):
            analysis_validation : dict[str, bool | list[str]] = Field(default={})
            skim_production : dict[str, bool | list[str]]  = Field(default={})
            run_quality : dict[str, bool | list[str]]  = Field(default={})


    In this model we make use of the @field_validator decorator which we use to validate each
    attribute that is not an empty dictionary. We are checking if the dictionary has of the following keys

    - any
    - by_group
    - by_name

    If none of these are found inside the non-empty dictionary then an assertion error is raised. How this is done is as follows:

    .. code-block:: python

        @field_validator('*')
        @classmethod
        def check_values_of_dictionaries(cls, modes_to_run_dict:dict) -> dict:
            if modes_to_run_dict:
                assert any(f for f in ['all', 'by_group', 'by_name'] if f in modes_to_run_dict.keys()), f'ModeType variant defined for modes_to_run in the config.toml does not contain a any of the following: all, by_name, by_group'
            return modes_to_run_dict

    Note in the field_validator decorator we use the '*', this is short-hand to tell pydantic every field of
    this model should be validated using this method.

    Note that this model does not allow any unknown fields. A ValidationError is raised if an unknown field is
    encountered. This would occur if a new ModeType is added to the ModeType Enum without reflecting that
    here in the model.

    """

    analysis_validation: dict[str, bool | list[str]] = Field(default={})
    skim_production: dict[str, bool | list[str]] = Field(default={})
    run_quality: dict[str, bool | list[str]] = Field(default={})

    model_config = ConfigDict(extra="forbid")



    @field_validator("*")
    @classmethod
    def check_values_of_dictionaries(cls, modes_to_run_dict: dict) -> dict:
        if modes_to_run_dict:
            assert any(
                f for f in ["all", "by_group", "by_name"] if f in modes_to_run_dict.keys()
            ), "ModeType variant defined for modes_to_run in the config.toml does not contain a any of the following: all, by_name, by_group"
        return modes_to_run_dict