Pydantic: Transform a value before it is assigned to a field?

user_14509475 0
Viewed 494

I have the following model

class Window(BaseModel):
    size: tuple[int, int]

and I would like to instantiate it like this:

fields = {'size': '1920x1080'}
window = Window(**fields)

Of course this fails since the value of 'size' is not of the correct type. However, I would like to add logic so that the value is split at x, i.e.:

def transform(raw: str) -> tuple[int, int]:
    x, y = raw.split('x')
    return int(x), int(y)

Does Pydantic support this?

pydantic python-3.x python
2 Answers

You can implement such a behaviour with pydantic's validator. Given your predefined function:

def transform(raw: str) -> tuple[int, int]:
    x, y = raw.split('x')
    return int(x), int(y)

You can implement it in your class like this:

from pydantic import BaseModel, validator


class Window(BaseModel):
    
    size: tuple[int, int]
    _extract_size = validator('size', pre=True, allow_reuse=True)(transform)

Note the pre=True argument passed to the validator. It means that it will be run before the default validator that checks if size is a tuple.

Now:

fields = {'size': '1920x1080'}
window = Window(**fields)
print(window)
# output: size=(1920, 1080)

Note that after that, you won't be able to instantiate your Window with a tuple for size.

fields2 = {'size': (800, 600)}
window2 = Window(**fields2)
# AttributeError: 'tuple' object has no attribute 'split'

In order to overcome that, you could simply bypass the function if a tuple is passed by altering slightly your code:

def transform(raw: str) -> tuple[int, int]:
    if type(raw) == tuple:
        return raw
    x, y = raw.split('x')
    return int(x), int(y)

Which should give:

fields2 = {'size': (800, 600)}
window2 = Window(**fields2)
print(window2)
# output: size:(800, 600)
user_126627450

Just to share a convtools-based alternative:

from convtools.contrib.models import (
    DictModel,
    build,
    casters,
    validate,
    validators,
)

class Window(DictModel):
    size: tuple[int, int] = (
        validate(validators.Type(str))
        .cast(casters.CustomUnsafe(lambda s: s.split("x")))
        .cast()
    )

obj, errors = build(Window, {"size": "1920x1080"})
# In [12]: obj
# Out[12]: Window(size=(1920, 1080))

obj, errors = build(Window, {"size": "1920 1080"})
# In [14]: errors
# Out[14]: {'size': {'__ERRORS': {'length': 'not enough values to unpack (expected 2, got 1)'}}}

Or if a custom method is desired:

class Window(DictModel):
    size: tuple[int, int] = field(cls_method=True)

    @classmethod
    def get_size(cls, data):
        try:
            x, y = data["size"].split("x")
            return (int(x), int(y)), None
        except Exception:  # this is bad, but this is only an example
            return None, {"invalid_size": True}

Docs: https://convtools.readthedocs.io/en/latest/models.html

Github: https://github.com/westandskif/convtools

user_53877380
Related