Shared via mypy Playground

main.py

from __future__ import annotations

import datetime
import functools
import inspect
from typing import (
    TYPE_CHECKING,
    Any,
    Concatenate,
    Generic,
    ParamSpec,
    Protocol,
    TypeVar,
    reveal_type,
)

if TYPE_CHECKING:
    from collections.abc import Callable


# Mock cache class for demo
class Cache:
    def get(self, key: str, default: Any | None = None) -> Any: ...
    def set(self, key: str, value: Any, timeout: int | None = None) -> None: ...


cache = Cache()


_P = ParamSpec('_P')
_T_co = TypeVar('_T_co', covariant=True)


class _Never:
    pass


_never = _Never()

_DefaultCacheTime = int(datetime.timedelta(days=1).total_seconds())  # 24 hours before cached result expires


class MethodWithCache(Protocol, Generic[_P, _T_co]):
    def __call__(
        self, dummy: _Never = _never, /, use_cache: bool = False, *args: _P.args, **kwargs: _P.kwargs
    ) -> _T_co: ...


def _with_redis_cache(
    f: Callable[Concatenate[API, _P], _T_co],
) -> MethodWithCache[_P, _T_co]:
    func_name = f.__name__
    f_sig = inspect.signature(f)
    cache_args = [p for p in f_sig.parameters if p != 'self']

    @functools.wraps(f)
    def wrapped_f(*args: _P.args, use_cache: bool = False, **kwargs: _P.kwargs) -> _T_co:
        cache_key = ''
        if use_cache:
            bound_args = f_sig.bind(*args, **kwargs)
            bound_args.apply_defaults()
            arguments = bound_args.arguments
            self = arguments['self']
            cache_attrs = ",".join(f'{p}={arguments[p]}' for p in cache_args)
            cache_key = f'{self}.{func_name}({cache_attrs})'

            cached_response: _T_co | None = cache.get(cache_key)
            if cached_response:
                return cached_response

        result = f(*args, **kwargs)

        if use_cache:
            cache.set(cache_key, result, _DefaultCacheTime)

        return result

    return wrapped_f


class API:
    @_with_redis_cache
    def get_campaign_message(
        self, message_type: str, message_id: str | None = None, include_template: bool = False
    ) -> dict[str, str]:
        return {}


api = API()
reveal_type(api.get_campaign_message)
# These should not give an error
api.get_campaign_message(message_type='type')
api.get_campaign_message(message_type='type', message_id='message_id')
api.get_campaign_message(message_type='type', message_id='message_id', include_template=False)
api.get_campaign_message(message_type='type', message_id='message_id', include_template=True)
api.get_campaign_message(message_type='type', message_id='message_id', use_cache=False, include_template=True)
api.get_campaign_message(message_type='type', use_cache=True, include_template=True)

# These should give an error
api.get_campaign_message()  # message type is required
api.get_campaign_message('type')  # message_type must be passed as a kwarg
api.get_campaign_message(False, message_type='type')  # false for use_cache as positional
api.get_campaign_message(message_type='type', message_id=123)  # int for message_id
api.get_campaign_message(message_type='type', message_id='message_id', include_template2=False)  # invalid kwarg
api.get_campaign_message(message_type='type', message_id='message_id', use_cache2=False)  # invalid kwarg
api.get_campaign_message(message_type='type', message_id='message_id', include_template=1)  # invalid arg type
api.get_campaign_message(message_type='type', message_id='message_id', use_cache=1)  # invalid arg type