Skip to content

Commit e5e462b

Browse files
Refactor @reawaitable decorator (#2029)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 0d7c37f commit e5e462b

File tree

1 file changed

+16
-6
lines changed

1 file changed

+16
-6
lines changed

returns/primitives/reawaitable.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
from collections.abc import Awaitable, Callable, Generator
2-
from typing import NewType, TypeVar, cast, final
2+
from functools import wraps
3+
from typing import NewType, ParamSpec, TypeVar, cast, final
34

45
_ValueType = TypeVar('_ValueType')
5-
_FunctionCoroType = TypeVar('_FunctionCoroType', bound=Callable[..., Awaitable])
6+
_AwaitableT = TypeVar('_AwaitableT', bound=Awaitable)
7+
_Ps = ParamSpec('_Ps')
68

79
_Sentinel = NewType('_Sentinel', object)
810
_sentinel: _Sentinel = cast(_Sentinel, object())
@@ -104,7 +106,9 @@ async def _awaitable(self) -> _ValueType:
104106
return self._cache # type: ignore
105107

106108

107-
def reawaitable(coro: _FunctionCoroType) -> _FunctionCoroType:
109+
def reawaitable(
110+
coro: Callable[_Ps, _AwaitableT],
111+
) -> Callable[_Ps, _AwaitableT]:
108112
"""
109113
Allows to decorate coroutine functions to be awaitable multiple times.
110114
@@ -124,6 +128,12 @@ def reawaitable(coro: _FunctionCoroType) -> _FunctionCoroType:
124128
>>> assert anyio.run(main) == 3
125129
126130
"""
127-
return lambda *args, **kwargs: ReAwaitable( # type: ignore
128-
coro(*args, **kwargs),
129-
)
131+
132+
@wraps(coro)
133+
def decorator(
134+
*args: _Ps.args,
135+
**kwargs: _Ps.kwargs,
136+
) -> _AwaitableT:
137+
return ReAwaitable(coro(*args, **kwargs)) # type: ignore[return-value]
138+
139+
return decorator

0 commit comments

Comments
 (0)