Data conversion on multiple fields with Pydantic

[laurent-laporte-pro]

Overall, this code example showcases the use of Pydantic to define a data model with validation and conversion logic, ensuring that the input values are in the expected format and can be seamlessly converted between different representations.

The Colors class bellow represents a color object that can be initialized with RGB components (color_r, color_g, color_b) or a CSS-like color string (color). The purpose of the class is to validate the input values and convert between the different representations.

Here's a breakdown of the key elements in the code:

1. The Colors class inherits from BaseModel.

2. The class defines four attributes: color_r, color_g, color_b, and color. All of these attributes are optional, meaning they can be None.

3. The Colors class has an inner class called Config, which specifies configuration options for the model. In this case, alias_generator is set to to_camel_case, which means the attribute aliases should be generated in camelCase. extra is set to Extra.forbid, which ensures that the model doesn't allow any extra fields that are not defined in the class. allow_population_by_field_name is set to True, enabling the population of fields using snake_case.

4. The validate_colors method is a root validator defined using the @root_validator decorator. It takes the values dictionary as input, representing the attribute values to be validated. This method performs the conversion between CSS colors and RGB components in both directions.

from typing import Optional, Dict, Any

from pydantic import BaseModel, root_validator, Extra

from antarest.core.utils.string import to_camel_case


class Colors(BaseModel):
    """
    Example of Colors class that support RGB component or CSS-like color.
    """
    color_r: Optional[int]
    color_g: Optional[int]
    color_b: Optional[int]
    color: Optional[str]

    class Config:
        alias_generator = to_camel_case
        extra = Extra.forbid
        allow_population_by_field_name = True

    @root_validator(pre=True)
    def validate_colors(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """
        This validator converts CSS colors into RGB components, in both directions.
        The conversion considers aliases in camelCase and field names in snake_case.
        Note that the validation is only effective if the object is initialized in
        the conventional way. If you use `construct`, the validation is not called.

        Args:
            values: The dictionary of attribute values to validate.

        Returns:
            The validated dictionary of attribute values.

        Raises:
            ValueError: If the validation fails (this exception is then converted into a `ValidationError` by Pydantic).
        """
        color = values.get("color")
        if color is not None:
            r, g, b = map(int, color.split(","))
            values["colorR"] = r
            values["colorG"] = g
            values["colorB"] = b
            return values

        r = values.get("colorR")
        g = values.get("colorG")
        b = values.get("colorB")
        if r is not None and g is not None and b is not None:
            values["color"] = ",".join(map(str, (r, g, b)))
            return  values

        if cls.Config.allow_population_by_field_name:
            r = values.get("color_r")
            g = values.get("color_g")
            b = values.get("color_b")
            if r is not None and g is not None and b is not None:
                values["color"] = ",".join(map(str, (r, g, b)))
                return values

        return values

The test_colors function demonstrates various test cases for validating the Colors class. It checks if the validation and conversion work as expected for different input scenarios.

def test_colors():
    # Validation is not done with `construct`.
    k = Colors.construct(color="255,128,0")
    assert k.dict() == {
        "color_r": None,
        "color_g": None,
        "color_b": None,
        "color": "255,128,0",
    }

    # Test validation with CSS-like color.
    k = Colors(color="255,128,0")
    assert k.dict() == {
        "color_r": 255,
        "color_g": 128,
        "color_b": 0,
        "color": "255,128,0",
    }

    # Test validation with field names in snake_case.
    if Colors.Config.allow_population_by_field_name:
        k = Colors(color_r=255, color_g=128, color_b=0)
        assert k.dict() == {
            "color_r": 255,
            "color_g": 128,
            "color_b": 0,
            "color": "255,128,0",
        }

    # Test validation with field names in camelCase.
    # noinspection Pydantic,PyArgumentList
    k = Colors(colorR=255, colorG=128, colorB=0)
    assert k.dict() == {
        "color_r": 255,
        "color_g": 128,
        "color_b": 0,
        "color": "255,128,0",
    }