From 9fd466b83872d1a19a403974ffc8fb8064a26616 Mon Sep 17 00:00:00 2001
From: Daniil Fajnberg <mail@daniil.fajnberg.de>
Date: Sat, 27 Nov 2021 12:30:48 +0100
Subject: [PATCH] async session decorator

---
 src/webutils/__init__.py |  1 +
 src/webutils/util.py     | 55 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 src/webutils/__init__.py
 create mode 100644 src/webutils/util.py

diff --git a/src/webutils/__init__.py b/src/webutils/__init__.py
new file mode 100644
index 0000000..b647f66
--- /dev/null
+++ b/src/webutils/__init__.py
@@ -0,0 +1 @@
+from .util import in_async_session
diff --git a/src/webutils/util.py b/src/webutils/util.py
new file mode 100644
index 0000000..af336a3
--- /dev/null
+++ b/src/webutils/util.py
@@ -0,0 +1,55 @@
+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)