Traits in Python with Protocol and Pydantic

traits_demo.py

from typing import Protocol, runtime_checkable
from pydantic import BaseModel, BeforeValidator
import inspect
from functools import wraps
from typing import Annotated, Optional, get_type_hints, Protocol, runtime_checkable

def enforce_protocol(func):
    sig = inspect.signature(func)
    hints = get_type_hints(func)
    @wraps(func)
    def wrapper(*args, **kwargs):
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()
        for name, value in bound.arguments.items():
            expected_type = hints.get(name)
            if expected_type and isinstance(expected_type, type) and issubclass(expected_type, Protocol):
                if not isinstance(value, expected_type):
                    raise TypeError(f"Argument '{name}' = {value!r} does not satisfy protocol {expected_type.__name__}")
        return func(*args, **kwargs)
    return wrapper

def validator(s: str) -> str:
    if len(s) == 1:
        return s * 10
    return s

@runtime_checkable
class Proto(Protocol):
    foo: str


class A(BaseModel):
    foo: str
    bar: int
    kek: Annotated[Optional[str], BeforeValidator(validator)] = "kek"

class B(BaseModel):
    bar: int

a = A(foo="hello", bar=42, kek="a")
b = B(bar=2)


@enforce_protocol
def f(s:str, c: Proto, i:int):
    print(f"{c.foo=}, {c.kek=}")

f("s", a, 1)
f("s", b, 1)  # fails

To be mixed with pydantic_partial