From b9f3ff1609232492490f4cfb73c556c255ba88ee Mon Sep 17 00:00:00 2001
From: Gabe Ruttner <gabriel.ruttner@gmail.com>
Date: Wed, 6 Mar 2024 15:24:44 -0800
Subject: [PATCH] feat(py-sdk): async steps (#246)

* feat: async wrapper

* feat: async examples

* release: bump version 0.15.1
---
 python-sdk/examples/async/event_test.py | 13 +++++++++
 python-sdk/examples/async/worker.py     | 36 +++++++++++++++++++++++++
 python-sdk/hatchet_sdk/hatchet.py       | 20 +++++++++-----
 python-sdk/pyproject.toml               |  2 +-
 4 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 python-sdk/examples/async/event_test.py
 create mode 100644 python-sdk/examples/async/worker.py

diff --git a/python-sdk/examples/async/event_test.py b/python-sdk/examples/async/event_test.py
new file mode 100644
index 000000000..8b4def4e8
--- /dev/null
+++ b/python-sdk/examples/async/event_test.py
@@ -0,0 +1,13 @@
+from hatchet_sdk import new_client
+from dotenv import load_dotenv
+
+load_dotenv()
+
+client = new_client()
+
+client.event.push(
+    "user:create",
+    {
+        "test": "test"
+    }
+)
\ No newline at end of file
diff --git a/python-sdk/examples/async/worker.py b/python-sdk/examples/async/worker.py
new file mode 100644
index 000000000..f0dee860c
--- /dev/null
+++ b/python-sdk/examples/async/worker.py
@@ -0,0 +1,36 @@
+from hatchet_sdk import Hatchet, Context
+from dotenv import load_dotenv
+import asyncio
+
+load_dotenv()
+
+hatchet = Hatchet(debug=True)
+
+@hatchet.workflow(on_events=["user:create"])
+class AsyncWorkflow:
+    def __init__(self):
+        self.my_value = "test"
+
+    @hatchet.step(timeout='5s')
+    def step1(self, context: Context):
+        async def async_step1():
+            print("started step1")
+            await asyncio.sleep(2)
+            print("finished step1")
+            return 'result'
+
+        res = asyncio.run(async_step1())
+        print(res)
+        return {'test': 'test'}
+
+    @hatchet.step(parents=["step1"],timeout='4s')
+    async def step2(self, context):
+        print("started async step2")
+        await asyncio.sleep(2)
+        print("finished step2")
+
+workflow = AsyncWorkflow()
+worker = hatchet.worker('test-worker', max_runs=4)
+worker.register_workflow(workflow)
+
+worker.start()
\ No newline at end of file
diff --git a/python-sdk/hatchet_sdk/hatchet.py b/python-sdk/hatchet_sdk/hatchet.py
index 3928e6574..1f8f24cd3 100644
--- a/python-sdk/hatchet_sdk/hatchet.py
+++ b/python-sdk/hatchet_sdk/hatchet.py
@@ -1,5 +1,7 @@
 from .client import new_client 
 from typing import List
+import asyncio
+from functools import wraps
 from .workflow import WorkflowMeta
 from .worker import Worker
 from .logger import logger
@@ -38,14 +40,20 @@ class Hatchet:
         
         return inner
 
-    def step(self, name : str='', timeout : str='', parents : List[str] = [], retries : int = 0):
+    def step(self, name: str='', timeout: str='', parents: List[str] = [], retries: int = 0):
         def inner(func):
-            func._step_name = name or func.__name__
-            func._step_parents = parents
-            func._step_timeout = timeout
-            func._step_retries = retries
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                if asyncio.iscoroutinefunction(func):
+                    return asyncio.run(func(*args, **kwargs))
+                else:
+                    return func(*args, **kwargs)
 
-            return func
+            wrapper._step_name = name or func.__name__
+            wrapper._step_parents = parents
+            wrapper._step_timeout = timeout
+            wrapper._step_retries = retries
+            return wrapper
 
         return inner
     
diff --git a/python-sdk/pyproject.toml b/python-sdk/pyproject.toml
index 155e66122..3d7cd1eb8 100644
--- a/python-sdk/pyproject.toml
+++ b/python-sdk/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "0.15.0"
+version = "0.15.1"
 description = ""
 authors = ["Alexander Belanger <alexander@hatchet.run>"]
 readme = "README.md"