From 846166c27ab823b7b229f5523efcca8aaf115374 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kacper=20Kali=C5=84ski?=
 <47140412+KaQuMiQ@users.noreply.github.com>
Date: Tue, 19 Nov 2024 11:33:02 +0100
Subject: [PATCH] Add wrap_async helper

---
 src/haiway/__init__.py           |  2 ++
 src/haiway/helpers/__init__.py   |  3 ++-
 src/haiway/helpers/asynchrony.py | 17 +++++++++++++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/haiway/__init__.py b/src/haiway/__init__.py
index 56b7f3c..686b4eb 100644
--- a/src/haiway/__init__.py
+++ b/src/haiway/__init__.py
@@ -15,6 +15,7 @@
     throttle,
     timeout,
     traced,
+    wrap_async,
 )
 from haiway.state import State
 from haiway.types import (
@@ -76,4 +77,5 @@
     "timeout",
     "traced",
     "when_missing",
+    "wrap_async",
 ]
diff --git a/src/haiway/helpers/__init__.py b/src/haiway/helpers/__init__.py
index 2caf455..85e15e6 100644
--- a/src/haiway/helpers/__init__.py
+++ b/src/haiway/helpers/__init__.py
@@ -1,4 +1,4 @@
-from haiway.helpers.asynchrony import asynchronous
+from haiway.helpers.asynchrony import asynchronous, wrap_async
 from haiway.helpers.caching import cache
 from haiway.helpers.retries import retry
 from haiway.helpers.throttling import throttle
@@ -14,4 +14,5 @@
     "throttle",
     "timeout",
     "traced",
+    "wrap_async",
 ]
diff --git a/src/haiway/helpers/asynchrony.py b/src/haiway/helpers/asynchrony.py
index e4b77a5..97b4020 100644
--- a/src/haiway/helpers/asynchrony.py
+++ b/src/haiway/helpers/asynchrony.py
@@ -9,9 +9,26 @@
 
 __all__ = [
     "asynchronous",
+    "wrap_async",
 ]
 
 
+def wrap_async[**Args, Result](
+    function: Callable[Args, Coroutine[None, None, Result]] | Callable[Args, Result],
+    /,
+) -> Callable[Args, Coroutine[None, None, Result]]:
+    if iscoroutinefunction(function):
+        return function
+
+    else:
+
+        async def async_function(*args: Args.args, **kwargs: Args.kwargs) -> Result:
+            return cast(Callable[Args, Result], function)(*args, **kwargs)
+
+        _mimic_async(function, within=async_function)
+        return async_function
+
+
 @overload
 def asynchronous[**Args, Result]() -> (
     Callable[