Python decorator with TypeVar and @overload

decorator_overload.py

import datetime as dt
from dataclasses import dataclass
from typing import Callable, Generic, Optional, TypeVar, overload

INPUT_T = TypeVar("INPUT_T")
OUTPUT_T = TypeVar("OUTPUT_T")


class pass_thru_none(Generic[INPUT_T, OUTPUT_T]):
    def __init__(self, func: Callable[[INPUT_T], OUTPUT_T]) -> None:
        self.func = func

    @overload
    def __call__(self, value: INPUT_T) -> OUTPUT_T:
        ...

    @overload
    def __call__(self, value: None) -> Optional[OUTPUT_T]:
        ...

    def __call__(self, value: Optional[INPUT_T]) -> Optional[OUTPUT_T]:
        if value is None:
            return None
        return self.func(value)


@pass_thru_none
def format_datetime(value: dt.datetime) -> str:
    return value.isoformat()


@dataclass
class UserModel:

    name: str
    create_date: dt.datetime
    login_date: Optional[dt.datetime]


@dataclass
class UserSchema:

    name: str
    createDate: str
    loginDate: Optional[str]

    @staticmethod
    def from_model(model: UserModel) -> "UserSchema":
        return UserSchema(
            name=model.name,
            createDate=format_datetime(model.create_date),
            loginDate=format_datetime(model.login_date),
        )


def test() -> None:
    model = UserModel(
        name="jdoe",
        create_date=dt.datetime(2023, 1, 1, 0, 0),
        login_date=None,
    )
    assert UserSchema.from_model(model) == UserSchema(
        name="jdoe",
        createDate="2023-01-01T00:00:00",
        loginDate=None,
    )