Shared via mypy Playground

main.py

import functools 
from typing import (
    Callable,
    Generic,
    NewType,
    TypeVar,
    ParamSpec,
)

UserId = NewType('UserId', str)
T = TypeVar('T')
P = ParamSpec('P')


class Auth(Generic[P, T]):
    def __init__(self, wrapped: Callable[P, T]) -> None:
        functools.update_wrapper(self, wrapped)
        self.wrapped = wrapped
        self.token = self.login()
    
    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
        return self.wrapped(*args, **kwargs)
        
    def login(self) -> dict[str, str]:
        return {'accessToken': 'wow', 'tokenType': 'Bearer'}
        

def with_auth(wrapped: Callable[P, T]) -> Auth[P, T]:
    @Auth
    def decorated(*args: P.args, **kwargs: P.kwargs) -> T:
        return wrapped(*args, **kwargs)
    return decorated

class Client:
    @with_auth
    def one(self, user_id: UserId) -> str:
        return f'Making call with token {self.one.token}'
        
    @classmethod
    @Auth
    def two(cls, user_id: UserId) -> str:
        return f'Making call with token {cls.two.token}'
        
    @staticmethod
    @Auth
    def three(user_id: UserId) -> str:
        return f'Making call with token {Client.three.token}'

@Auth
def four(user_id: UserId) -> str:
    return f'Making call with token {four.token}'
    
@with_auth
def five(user_id: UserId) -> str:
    return f'Making call with token {five.token}'

# Succeeds as expected.
UserId('asdf')
# Fails Successfully: Argument 1 to "UserId" has incompatible type "int"; expected "str"  [arg-type]
UserId(42)

client = Client()

# Should succeed but doesn't: Missing positional argument "self" in call to "__call__" of "Auth"  [call-arg]
client.one(user_id=UserId('asdf'))
# These pass successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "str"; expected "UserId"
client.one(client, user_id=UserId('asdf'))
Client.one(Client(), user_id=UserId('asdf'))
# Fails for the wrong reason: Missing positional argument "self" in call to "__call__" of "Auth"  [call-arg] 
client.one(user_id='asdf')
# These fail for the correct reason: 
client.one(client, user_id='asdf')
Client.one(Client(), user_id='asdf')
# Fails for the wrong reason: Missing positional argument "self" in call to "__call__" of "Auth"  [call-arg]
client.one(user_id=42)
# These fail successfully
client.one(client, user_id=42)
Client.one(Client(), user_id=42)

# Should succeed but doesn't: Missing positional argument "cls" in call to "__call__" of "Auth"  [call-arg]
Client.two(user_id=UserId('asdf'))
# Fails for the wrong reason: Missing positional argument "cls" in call to "__call__" of "Auth"  [call-arg]
Client.two(user_id='asdf')
# Fails for the wrong reason: Missing positional argument "cls" in call to "__call__" of "Auth"  [call-arg]
Client.two(user_id=42)

# Succeeds as expected.
Client.three(user_id=UserId('asdf'))
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "str"; expected "UserId"  [arg-type]
Client.three(user_id='asdf')
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "int"; expected "UserId"  [arg-type]
Client.three(user_id=42)

# Succeeds as expected.
four(user_id=UserId('asdf'))
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "str"; expected "UserId"  [arg-type]
four(user_id='asdf')
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "int"; expected "UserId"  [arg-type]
four(user_id=42)

# Succeeds as expected
five(user_id=UserId('asdf'))
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "str"; expected "UserId"  [arg-type]
five(user_id='asdf')
# Fails Successfully: Argument "user_id" to "__call__" of "Auth" has incompatible type "int"; expected "UserId"  [arg-type]
five(user_id=42)