56 lines
2.9 KiB
Python
56 lines
2.9 KiB
Python
import logging
|
|
from functools import wraps
|
|
from typing import Callable, Dict, Any
|
|
|
|
from aiohttp.client import ClientSession
|
|
|
|
|
|
LOGGER_NAME = 'webutils'
|
|
logger = logging.getLogger(LOGGER_NAME)
|
|
|
|
|
|
def in_async_session(_func: Callable = None, *,
|
|
session_kwargs: Dict[str, Any] = None, session_param_name: str = 'session') -> Callable:
|
|
"""
|
|
Useful decorator for any async function that uses the `aiohttp.ClientSession` to make requests.
|
|
|
|
Using this decorator allows the decorated function to have an optional session parameter,
|
|
without the need to ensure proper initialization and closing of a session within the function itself.
|
|
|
|
The wrapper has no effect, if a session object is passed into the function call, but if no session is passed,
|
|
it initializes one, passes it into the function and ensures that it is closed in the end.
|
|
|
|
Args:
|
|
_func:
|
|
If this decorator is used *with any* arguments, this will always be the decorated function itself.
|
|
This is a trick to allow the decorator to be used with as well as without arguments, i.e. in the form
|
|
`@in_async_session` or `@in_async_session(...)`.
|
|
session_kwargs (optional):
|
|
If passed a dictionary, it will be unpacked and passed as keyword arguments into the `ClientSession`
|
|
constructor, if and only if the decorator actually handles session initialization/closing,
|
|
i.e. only when the function is called **without** passing a session object into it.
|
|
session_param_name (optional):
|
|
The name of the decorated function's parameter that should be passed the session object as an argument.
|
|
In case the decorated function's session parameter is named anything other than "session", that name should
|
|
be provided here.
|
|
"""
|
|
def decorator(function: Callable) -> Callable:
|
|
# Using `functools.wraps` to preserve information about the actual function being decorated
|
|
# More details: https://docs.python.org/3/library/functools.html#functools.wraps
|
|
@wraps(function)
|
|
async def wrapper(*args, **kwargs):
|
|
"""The actual function wrapper that may perform the session initialization and closing."""
|
|
temp_session = False
|
|
if not any(isinstance(arg, ClientSession) for arg in args) and kwargs.get(session_param_name) is None:
|
|
logger.debug("Starting temporary client session")
|
|
kwargs[session_param_name] = ClientSession(**session_kwargs if session_kwargs is not None else {})
|
|
temp_session = True
|
|
try:
|
|
return await function(*args, **kwargs)
|
|
finally:
|
|
if temp_session:
|
|
await kwargs[session_param_name].close()
|
|
logger.debug("Temporary client session closed")
|
|
return wrapper
|
|
return decorator if _func is None else decorator(_func)
|