Skip to content

Instantly share code, notes, and snippets.

@LefterisJP
Last active June 12, 2022 15:18
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save LefterisJP/d7e2304b73820aa1f796b6e5bb93a0ad to your computer and use it in GitHub Desktop.
Save LefterisJP/d7e2304b73820aa1f796b6e5bb93a0ad to your computer and use it in GitHub Desktop.
how_to_type_this
import random
from functools import wraps
from typing import TYPE_CHECKING, Any, Callable, Union
from typing_extensions import Concatenate, ParamSpec
P = ParamSpec('P')
def maybe_inject_int(method: Callable[Concatenate['DBHandler', int, P], Any]) -> Union[
Callable[Concatenate['DBHandler', P], Any],
Callable[Concatenate['DBHandler', int, P], Any],
]:
"""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
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=}'
db = DBHandler()
if TYPE_CHECKING:
reveal_type(db.foo)
print(db.foo(5, 'John'))
print(db.foo('Mary'))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment