feat: expose generator

python-sdk/examples/manual_trigger/trigger-generator.py

@@ -0,0 +1,16 @@
+from hatchet_sdk import Hatchet
+from dotenv import load_dotenv
+import json
+
+load_dotenv()
+
+client = Hatchet().client
+
+workflowRunId = client.admin.run_workflow("ManualTriggerWorkflow", {
+    "test": "test"
+})
+
+for event in client.listener.generator(workflowRunId):
+    print('EVENT: ' + event.type + ' ' + json.dumps(event.payload))
+
+# TODO - need to hangup the listener if the workflow is completed

python-sdk/hatchet_sdk/clients/listener.py

@@ -48,7 +48,7 @@ def __init__(self, client: DispatcherStub, token):
         self.client = client
         self.token = token
 
-    def on(self, workflowRunId: str, handler: callable):
+    def generator(self, workflowRunId: str) -> List[StepRunEvent]:
         listener = self.retry_subscribe(workflowRunId)
 
         while True:
@@ -67,7 +67,8 @@ def on(self, workflowRunId: str, handler: callable):
                         payload = json.loads(workflow_event.eventPayload)
 
                     # call the handler
-                    handler(StepRunEvent(type=eventType, payload=payload))
+                    event = StepRunEvent(type=eventType, payload=payload)
+                    yield event
 
             except grpc.RpcError as e:
                 # Handle different types of errors
@@ -86,6 +87,12 @@ def on(self, workflowRunId: str, handler: callable):
                     logger.error(f"Failed to receive message: {e}")
                     break
 
+    def on(self, workflowRunId: str, handler: callable = None):
+        for event in self.generator(workflowRunId):
+            # call the handler if provided
+            if handler:
+                handler(event)
+
     def retry_subscribe(self, workflowRunId: str):
         retries = 0