Merge pull request #2 from hatchet-dev/feat--typed-dicts

Remove typed dicts in favor of Pydantic

Author: mrkaye97

examples/affinity-workers/event.py

@@ -1,5 +1,6 @@
 from dotenv import load_dotenv
 
+from hatchet_sdk.clients.events import PushEventOptions
 from hatchet_sdk.hatchet import Hatchet
 
 load_dotenv()
@@ -9,5 +10,5 @@
 hatchet.event.push(
     "affinity:run",
     {"test": "test"},
-    options={"additional_metadata": {"hello": "moon"}},
+    options=PushEventOptions(additional_metadata={"hello": "moon"}),
 )

examples/affinity-workers/worker.py

@@ -1,6 +1,7 @@
 from dotenv import load_dotenv
 
 from hatchet_sdk import Context, Hatchet, WorkerLabelComparator
+from hatchet_sdk.labels import DesiredWorkerLabel
 
 load_dotenv()
 
@@ -11,12 +12,12 @@
 class AffinityWorkflow:
     @hatchet.step(
         desired_worker_labels={
-            "model": {"value": "fancy-ai-model-v2", "weight": 10},
-            "memory": {
-                "value": 256,
-                "required": True,
-                "comparator": WorkerLabelComparator.LESS_THAN,
-            },
+            "model": DesiredWorkerLabel(value="fancy-ai-model-v2", weight=10),
+            "memory": DesiredWorkerLabel(
+                value=256,
+                required=True,
+                comparator=WorkerLabelComparator.LESS_THAN,
+            ),
         },
     )
     async def step(self, context: Context) -> dict[str, str | None]:

examples/bulk_fanout/stream.py

@@ -31,7 +31,7 @@ async def main() -> None:
     workflowRun = hatchet.admin.run_workflow(
         "Parent",
         {"n": 2},
-        options={"additional_metadata": {streamKey: streamVal}},
+        options=TriggerWorkflowOptions(additional_metadata={streamKey: streamVal}),
     )
 
     # Stream all events for the additional meta key value

examples/bulk_fanout/worker.py

@@ -4,7 +4,7 @@
 from dotenv import load_dotenv
 
 from hatchet_sdk import Context, Hatchet
-from hatchet_sdk.clients.admin import ChildWorkflowRunDict
+from hatchet_sdk.clients.admin import ChildTriggerWorkflowOptions, ChildWorkflowRunDict
 
 load_dotenv()
 
@@ -22,18 +22,17 @@ async def spawn(self, context: Context) -> dict[str, list[Any]]:
 
         n = context.workflow_input().get("n", 100)
 
-        child_workflow_runs: list[ChildWorkflowRunDict] = []
-
-        for i in range(n):
-
-            child_workflow_runs.append(
-                {
-                    "workflow_name": "BulkChild",
-                    "input": {"a": str(i)},
-                    "key": f"child{i}",
-                    "options": {"additional_metadata": {"hello": "earth"}},
-                }
+        child_workflow_runs = [
+            ChildWorkflowRunDict(
+                workflow_name="BulkChild",
+                input={"a": str(i)},
+                key=f"child{i}",
+                options=ChildTriggerWorkflowOptions(
+                    additional_metadata={"hello": "earth"}
+                ),
             )
+            for i in range(n)
+        ]
 
         if len(child_workflow_runs) == 0:
             return {}

examples/dedupe/worker.py

@@ -4,7 +4,7 @@
 
 from dotenv import load_dotenv
 
-from hatchet_sdk import Context, Hatchet
+from hatchet_sdk import ChildTriggerWorkflowOptions, Context, Hatchet
 from hatchet_sdk.clients.admin import DedupeViolationErr
 from hatchet_sdk.loader import ClientConfig
 
@@ -29,7 +29,9 @@ async def spawn(self, context: Context) -> dict[str, list[Any]]:
                             "DedupeChild",
                             {"a": str(i)},
                             key=f"child{i}",
-                            options={"additional_metadata": {"dedupe": "test"}},
+                            options=ChildTriggerWorkflowOptions(
+                                additional_metadata={"dedupe": "test"}
+                            ),
                         )
                     ).result()
                 )

examples/durable_sticky_with_affinity/worker.py

@@ -3,7 +3,13 @@
 
 from dotenv import load_dotenv
 
-from hatchet_sdk import Context, StickyStrategy, WorkerLabelComparator
+from hatchet_sdk import (
+    ChildTriggerWorkflowOptions,
+    Context,
+    StickyStrategy,
+    WorkerLabelComparator,
+)
+from hatchet_sdk.labels import DesiredWorkerLabel
 from hatchet_sdk.v2.callable import DurableContext
 from hatchet_sdk.v2.hatchet import Hatchet
 
@@ -15,17 +21,17 @@
 @hatchet.durable(
     sticky=StickyStrategy.HARD,
     desired_worker_labels={
-        "running_workflow": {
-            "value": "True",
-            "required": True,
-            "comparator": WorkerLabelComparator.NOT_EQUAL,
-        },
+        "running_workflow": DesiredWorkerLabel(
+            value="True",
+            required=True,
+            comparator=WorkerLabelComparator.NOT_EQUAL,
+        ),
     },
 )
 async def my_durable_func(context: DurableContext) -> dict[str, Any]:
     try:
         ref = await context.aio.spawn_workflow(
-            "StickyChildWorkflow", {}, options={"sticky": True}
+            "StickyChildWorkflow", {}, options=ChildTriggerWorkflowOptions(sticky=True)
         )
         result = await ref.result()
     except Exception as e:
@@ -39,11 +45,11 @@ async def my_durable_func(context: DurableContext) -> dict[str, Any]:
 class StickyChildWorkflow:
     @hatchet.step(
         desired_worker_labels={
-            "running_workflow": {
-                "value": "True",
-                "required": True,
-                "comparator": WorkerLabelComparator.NOT_EQUAL,
-            },
+            "running_workflow": DesiredWorkerLabel(
+                value="True",
+                required=True,
+                comparator=WorkerLabelComparator.NOT_EQUAL,
+            ),
         },
     )
     async def child(self, context: Context) -> dict[str, str | None]:

examples/events/test_event.py

@@ -24,34 +24,34 @@ async def test_async_event_push(aiohatchet: Hatchet) -> None:
 @pytest.mark.asyncio(scope="session")
 async def test_async_event_bulk_push(aiohatchet: Hatchet) -> None:
 
-    events: List[BulkPushEventWithMetadata] = [
-        {
-            "key": "event1",
-            "payload": {"message": "This is event 1"},
-            "additional_metadata": {"source": "test", "user_id": "user123"},
-        },
-        {
-            "key": "event2",
-            "payload": {"message": "This is event 2"},
-            "additional_metadata": {"source": "test", "user_id": "user456"},
-        },
-        {
-            "key": "event3",
-            "payload": {"message": "This is event 3"},
-            "additional_metadata": {"source": "test", "user_id": "user789"},
-        },
+    events = [
+        BulkPushEventWithMetadata(
+            key="event1",
+            payload={"message": "This is event 1"},
+            additional_metadata={"source": "test", "user_id": "user123"},
+        ),
+        BulkPushEventWithMetadata(
+            key="event2",
+            payload={"message": "This is event 2"},
+            additional_metadata={"source": "test", "user_id": "user456"},
+        ),
+        BulkPushEventWithMetadata(
+            key="event3",
+            payload={"message": "This is event 3"},
+            additional_metadata={"source": "test", "user_id": "user789"},
+        ),
     ]
-    opts: BulkPushEventOptions = {"namespace": "bulk-test"}
+    opts = BulkPushEventOptions(namespace="bulk-test")
 
     e = await aiohatchet.event.async_bulk_push(events, opts)
 
     assert len(e) == 3
 
     # Sort both lists of events by their key to ensure comparison order
-    sorted_events = sorted(events, key=lambda x: x["key"])
+    sorted_events = sorted(events, key=lambda x: x.key)
     sorted_returned_events = sorted(e, key=lambda x: x.key)
     namespace = "bulk-test"
 
     # Check that the returned events match the original events
     for original_event, returned_event in zip(sorted_events, sorted_returned_events):
-        assert returned_event.key == namespace + original_event["key"]
+        assert returned_event.key == namespace + original_event.key

examples/fanout/stream.py

@@ -31,7 +31,7 @@ async def main() -> None:
     workflowRun = hatchet.admin.run_workflow(
         "Parent",
         {"n": 2},
-        options={"additional_metadata": {streamKey: streamVal}},
+        options=TriggerWorkflowOptions(additional_metadata={streamKey: streamVal}),
     )
 
     # Stream all events for the additional meta key value

examples/fanout/sync_stream.py

@@ -31,7 +31,7 @@ def main() -> None:
     workflowRun = hatchet.admin.run_workflow(
         "Parent",
         {"n": 2},
-        options={"additional_metadata": {streamKey: streamVal}},
+        options=TriggerWorkflowOptions(additional_metadata={streamKey: streamVal}),
     )
 
     # Stream all events for the additional meta key value

examples/fanout/worker.py

@@ -3,7 +3,7 @@
 
 from dotenv import load_dotenv
 
-from hatchet_sdk import Context, Hatchet
+from hatchet_sdk import ChildTriggerWorkflowOptions, Context, Hatchet
 
 load_dotenv()
 
@@ -28,7 +28,9 @@ async def spawn(self, context: Context) -> dict[str, Any]:
                         "Child",
                         {"a": str(i)},
                         key=f"child{i}",
-                        options={"additional_metadata": {"hello": "earth"}},
+                        options=ChildTriggerWorkflowOptions(
+                            additional_metadata={"hello": "earth"}
+                        ),
                     )
                 ).result()
             )