unpack

Author: timl3136

cadence/worker/_registry.py

@@ -7,25 +7,22 @@
 """
 
 import logging
-from typing import Callable, Dict, Optional
-from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Unpack, TypedDict
 
 
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class RegisterWorkflowOptions:
+class RegisterWorkflowOptions(TypedDict, total=False):
     """Options for registering a workflow."""
-    name: Optional[str] = None
-    alias: Optional[str] = None
+    name: Optional[str]
+    alias: Optional[str]
 
 
-@dataclass
-class RegisterActivityOptions:
+class RegisterActivityOptions(TypedDict, total=False):
     """Options for registering an activity."""
-    name: Optional[str] = None
-    alias: Optional[str] = None
+    name: Optional[str]
+    alias: Optional[str]
 
 
 class Registry:
@@ -46,7 +43,7 @@ def __init__(self):
     def workflow(
         self,
         func: Optional[Callable] = None,
-        **kwargs: RegisterWorkflowOptions
+        **kwargs: Unpack[RegisterWorkflowOptions]
     ) -> Callable:
         """
         Register a workflow function.
@@ -66,18 +63,19 @@ def workflow(
         options = RegisterWorkflowOptions(**kwargs)
 
         def decorator(f: Callable) -> Callable:
-            workflow_name = options.name or f.__name__
+            workflow_name = options.get('name') or f.__name__
 
             if workflow_name in self._workflows:
                 raise KeyError(f"Workflow '{workflow_name}' is already registered")
 
             self._workflows[workflow_name] = f
 
             # Register alias if provided
-            if options.alias:
-                if options.alias in self._workflow_aliases:
-                    raise KeyError(f"Workflow alias '{options.alias}' is already registered")
-                self._workflow_aliases[options.alias] = workflow_name
+            alias = options.get('alias')
+            if alias:
+                if alias in self._workflow_aliases:
+                    raise KeyError(f"Workflow alias '{alias}' is already registered")
+                self._workflow_aliases[alias] = workflow_name
 
             logger.info(f"Registered workflow '{workflow_name}'")
             return f
@@ -89,7 +87,7 @@ def decorator(f: Callable) -> Callable:
     def activity(
         self,
         func: Optional[Callable] = None,
-        **kwargs: RegisterActivityOptions
+        **kwargs: Unpack[RegisterActivityOptions]
     ) -> Callable:
         """
         Register an activity function.
@@ -109,18 +107,19 @@ def activity(
         options = RegisterActivityOptions(**kwargs)
 
         def decorator(f: Callable) -> Callable:
-            activity_name = options.name or f.__name__
+            activity_name = options.get('name') or f.__name__
 
             if activity_name in self._activities:
                 raise KeyError(f"Activity '{activity_name}' is already registered")
 
             self._activities[activity_name] = f
 
             # Register alias if provided
-            if options.alias:
-                if options.alias in self._activity_aliases:
-                    raise KeyError(f"Activity alias '{options.alias}' is already registered")
-                self._activity_aliases[options.alias] = activity_name
+            alias = options.get('alias')
+            if alias:
+                if alias in self._activity_aliases:
+                    raise KeyError(f"Activity alias '{alias}' is already registered")
+                self._activity_aliases[alias] = activity_name
 
             logger.info(f"Registered activity '{activity_name}'")
             return f

tests/cadence/worker/test_registry.py

@@ -170,7 +170,7 @@ def test_options_class(self, registration_type):
         if registration_type == "workflow":
             options = RegisterWorkflowOptions(name="custom_name", alias="custom_alias")
 
-            @reg.workflow(**options.__dict__)
+            @reg.workflow(**options)
             def test_func():
                 return "test"
 
@@ -179,7 +179,7 @@ def test_func():
         else:
             options = RegisterActivityOptions(name="custom_name", alias="custom_alias")
 
-            @reg.activity(**options.__dict__)
+            @reg.activity(**options)
             def test_func():
                 return "test"