"""
Provides a Golang-like ``defer()`` API using decorators, which allows grouping
resource initialization and cleanup in one place without extra indentations.
Example:
.. code-block:: python3
async def init(x):
...
async def cleanup(x):
...
@aiotools.adefer
async def do(defer): # <-- be aware of defer argument!
x = SomeResource()
await init(x)
defer(cleanup(x))
...
...
This is equivalent to:
.. code-block:: python3
async def do():
x = SomeResource()
await init(x)
try:
...
...
finally:
await cleanup(x)
Note that :class:`aiotools.context.AsyncContextGroup` or
:class:`contextlib.AsyncExitStack` serves well for the same purpose, but for simple
cleanups, this defer API makes your codes simple because it steps aside the main
execution context without extra indentations.
.. warning::
Any exception in the deferred functions is raised transparently, and may block
execution of the remaining deferred functions.
This behavior may be changed in the future versions, though.
"""
from __future__ import annotations
import functools
import inspect
from collections import deque
from collections.abc import Awaitable, Callable
from typing import Any, Concatenate, ParamSpec, TypeAlias, TypeVar
P = ParamSpec("P")
R = TypeVar("R")
__all__ = (
"defer",
"adefer",
"DeferFunc",
"AsyncDeferrable",
"AsyncDeferFunc",
)
DeferFunc: TypeAlias = Callable[[Callable[[], None]], None]
AsyncDeferrable: TypeAlias = (
Callable[[], Awaitable[Any]] | Callable[[], None] | Awaitable[None]
)
AsyncDeferFunc: TypeAlias = Callable[[AsyncDeferrable], None]
[docs]
def defer(func: Callable[Concatenate[DeferFunc, P], R]) -> Callable[P, R]:
"""
A synchronous version of the defer API.
It can only defer normal functions.
"""
assert not inspect.iscoroutinefunction(func), (
"the decorated function must not be async"
)
@functools.wraps(func)
def _wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
deferreds: deque[Callable[[], None]] = deque()
# the registration function passed to the original function
def defer(f: Callable[[], None]) -> None:
assert not inspect.iscoroutinefunction(f), (
"the deferred function must not be async"
)
assert not inspect.iscoroutine(f), (
"the deferred object must not be a coroutine"
)
deferreds.append(f)
try:
return func(defer, *args, **kwargs)
finally:
while deferreds:
f = deferreds.pop()
f()
return _wrapped
[docs]
def adefer(
func: Callable[Concatenate[AsyncDeferFunc, P], Awaitable[R]],
) -> Callable[P, Awaitable[R]]:
"""
An asynchronous version of the defer API.
It can defer coroutine functions, coroutines, and normal functions.
"""
assert inspect.iscoroutinefunction(func), "the decorated function must be async"
@functools.wraps(func)
async def _wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
deferreds: deque[AsyncDeferrable] = deque()
# the registration function passed to the original function
def defer(f: AsyncDeferrable) -> None:
deferreds.append(f)
try:
return await func(defer, *args, **kwargs)
finally:
while deferreds:
f = deferreds.pop()
if inspect.iscoroutinefunction(f):
await f()
elif inspect.iscoroutine(f):
await f
else:
f() # type: ignore[operator]
return _wrapped