this type checks!! yay

it_works.py

import random
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Dict, Protocol, Union, overload

from typing_extensions import Concatenate, ParamSpec

P = ParamSpec('P')


class MaybeInjectedInt(Protocol[P]):
    @overload
    def __call__(self, number: int, *args: P.args, **kwargs: P.kwargs) -> Any:
        ...

    @overload
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any:
        ...


def maybe_inject_int(method: Callable[Concatenate['DBHandler', int, P], Any]) -> MaybeInjectedInt[P]:
    """Wraps the method of DBHandler to maybe inject a positional int argument
    """
    @wraps(method)
    def _impl(self: 'DBHandler', *args: Any, **kwargs: Any) -> Union[
            Callable[P, Any],
            Callable[Concatenate[int, P], Any],
    ]:
        if isinstance(args[0], int):  # first argument is the int
            return method(self, *args, **kwargs)
        # else we need to inject the argument
        new_int = random.randint(1, 99)
        result = method(self, new_int, *args, **kwargs)
        return result
    return _impl  # type: ignore


class DBHandler:

    def __init__(self) -> None:
        self.id = 1

    @maybe_inject_int
    def foo(self, number: int, name: str) -> str:
        return f'{self.id=} {name=} {number=}'

    @maybe_inject_int
    def boo(self, number: int, data: Dict) -> Dict:
        if number > 0:
            return data
        else:
            return {}
        


db = DBHandler()

if TYPE_CHECKING:
    reveal_type(db.foo)

print(db.foo(5, 'John'))
print(db.foo('Mary'))

# Both commented ones correctly fail mypy check!
# db.foo(5, 5)
# db.foo(5)

if TYPE_CHECKING:
    reveal_type(db.boo)
db.boo(5, {'a': 1, 'b': 2})
db.boo({'a': 1, 'b': 2})

# Both commented ones correctly fail mypy check!
# db.boo(5, 'kek')
# db.boo(5)