diff --git a/src/wrapt/decorators.py b/src/wrapt/decorators.py index 41459b7..926ddaa 100644 --- a/src/wrapt/decorators.py +++ b/src/wrapt/decorators.py @@ -6,6 +6,7 @@ from functools import partial from inspect import isclass, signature from threading import Lock, RLock +from typing import Any, Callable, Optional, Type, TypeVar, overload from .arguments import formatargspec @@ -15,6 +16,10 @@ CallableObjectProxy, ) +F = TypeVar('F', bound=Callable[..., Any]) +A = TypeVar('A', bound=Callable[..., Any]) +T = TypeVar('T', bound=Any) + # Adapter wrapper for the wrapped function which will overlay certain # properties from the adapter function onto the wrapped function so that # functions such as inspect.getfullargspec(), inspect.signature() and @@ -141,7 +146,30 @@ def __call__(self, wrapped): # original wrapped function. -def decorator(wrapper=None, enabled=None, adapter=None, proxy=FunctionWrapper): +@overload +def decorator( + wrapper: None = None, + enabled: Optional[bool] = None, + adapter: Optional[A] = None, + proxy: Type[FunctionWrapper] = FunctionWrapper, +) -> partial[F]: ... + + +@overload +def decorator( + wrapper: F, + enabled: Optional[bool] = None, + adapter: Optional[A] = None, + proxy: Type[FunctionWrapper] = FunctionWrapper, +) -> F: ... + + +def decorator( + wrapper: Optional[F] = None, + enabled: Optional[bool] = None, + adapter: Optional[A] = None, + proxy: Type[FunctionWrapper] = FunctionWrapper, +) -> F | partial[F]: # The decorator should be supplied with a single positional argument # which is the wrapper function to be used to implement the # decorator. This may be preceded by a step whereby the keyword diff --git a/src/wrapt/wrappers.py b/src/wrapt/wrappers.py index 1631375..5c4c3cc 100644 --- a/src/wrapt/wrappers.py +++ b/src/wrapt/wrappers.py @@ -1,6 +1,10 @@ -import sys -import operator import inspect +import operator +import sys +from typing import Any, Generic, TypeVar + + +T = TypeVar("T", bound=Any) class _ObjectProxyMethods(object): @@ -61,10 +65,11 @@ def __new__(cls, name, bases, dictionary): class _ObjectProxyBase(metaclass=_ObjectProxyMetaType): ... -class ObjectProxy(_ObjectProxyBase): +class ObjectProxy(_ObjectProxyBase, Generic[T]): __slots__ = '__wrapped__' + __wrapped__: T - def __init__(self, wrapped): + def __init__(self, wrapped: T): object.__setattr__(self, '__wrapped__', wrapped) # Python 3.2+ has the __qualname__ attribute, but it does not