Pydantic Decorators V2

pydantic_decorators_v2.py

# PROS: Easy to use in CRUDs
# CONS: Invalid autocompletion

from inspect import isclass

from pydantic import BaseModel, create_model
from pydantic_core import SchemaSerializer, SchemaValidator


def omit(*fields):
    """Omit pydantic fields from model"""

    def dec(_class: BaseModel):
        for field in fields:
            _class.model_fields.pop(field, None)
        _clone = create_model(
            _class.__name__,
            __config__=_class.model_config,
            **{k: (v.annotation, v) for k, v in _class.model_fields.items()},
        )
        setattr(_clone, "__pydantic_parent_namespace__", {})
        return _clone

    return dec


def partial(*fields):
    """Turn pydantic fields into optional"""

    def dec(_class: BaseModel):
        _core = _class.__pydantic_core_schema__

        if _core["schema"]["type"] == "model-fields":
            _fields = _core["schema"]["fields"]
        if _core["schema"]["type"] == "model":
            _fields = _core["schema"]["schema"]["fields"]

        for field in fields:
            if _class.model_fields[field].is_required():
                _fields[field]["schema"] = {
                    "type": "default",
                    "schema": _fields[field]["schema"],
                    "default": None,
                }
                _class.model_fields[field].default = None

        _class.__pydantic_validator__ = SchemaValidator(_core)
        _class.__pydantic_serializer__ = SchemaSerializer(_core)

        return _class

    if fields and isclass(fields[0]) and issubclass(fields[0], BaseModel):
        _class = fields[0]
        fields = _class.model_fields
        return dec(_class)

    return dec


class Grandparent(BaseModel):
    name: str
    money: int
    benefit: str


@omit("benefit")
@partial("money")
class Parent(Grandparent):
    company: str


@omit("money")
@omit("company")
class Child(Parent):
    kindergarten: str


print(Child.model_fields)
print(Parent.model_fields)
print(Grandparent.model_fields)

Child(name="Artur", kindergarten="K")
Parent(name="Arnaldo", company="Company")
Grandparent(name="Armando", benefit="Medicare", money=0)