Redefine the workflow type to class instead of callable

Signed-off-by: Tim Li <[email protected]>

Author: timl3136

cadence/__init__.py

@@ -6,9 +6,13 @@
 
 # Import main client functionality
 from .client import Client
+from .worker import Registry
+from .workflow import workflow
 
 __version__ = "0.1.0"
 
 __all__ = [
     "Client",
+    "Registry",
+    "workflow",
 ]

cadence/worker/_registry.py

@@ -7,7 +7,7 @@
 """
 
 import logging
-from typing import Callable, Dict, Optional, Unpack, TypedDict, Sequence, overload
+from typing import Callable, Dict, Optional, Unpack, TypedDict, Sequence, overload, Type
 from cadence.activity import ActivityDefinitionOptions, ActivityDefinition, ActivityDecorator, P, T
 from cadence.workflow import WorkflowDefinition, WorkflowDefinitionOptions
 
@@ -35,50 +35,52 @@ def __init__(self) -> None:
 
     def workflow(
         self,
-        func: Optional[Callable] = None,
+        cls: Optional[Type] = None,
         **kwargs: Unpack[RegisterWorkflowOptions]
-    ) -> Callable:
+    ) -> Type:
         """
-        Register a workflow function.
-        
+        Register a workflow class.
+
         This method can be used as a decorator or called directly.
-        
+        Only supports class-based workflows.
+
         Args:
-            func: The workflow function to register
+            cls: The workflow class to register
             **kwargs: Options for registration (name, alias)
-            
+
         Returns:
-            The decorated function or the function itself
-            
+            The decorated class
+
         Raises:
             KeyError: If workflow name already exists
+            ValueError: If class workflow is invalid
         """
         options = RegisterWorkflowOptions(**kwargs)
-        
-        def decorator(f: Callable) -> Callable:
-            workflow_name = options.get('name') or f.__name__
-            
+
+        def decorator(target: Type) -> Type:
+            workflow_name = options.get('name') or target.__name__
+
             if workflow_name in self._workflows:
                 raise KeyError(f"Workflow '{workflow_name}' is already registered")
-            
+
             # Create WorkflowDefinition with type information
             workflow_opts = WorkflowDefinitionOptions(name=workflow_name)
-            workflow_def = WorkflowDefinition.wrap(f, workflow_opts)
+            workflow_def = WorkflowDefinition.wrap(target, workflow_opts)
             self._workflows[workflow_name] = workflow_def
-            
+
             # Register alias if provided
             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
-        
-        if func is None:
+            return target
+
+        if cls is None:
             return decorator
-        return decorator(func)
+        return decorator(cls)
 
     @overload
     def activity(self, func: Callable[P, T]) -> ActivityDefinition[P, T]:
@@ -142,22 +144,22 @@ def _register_activity(self, defn: ActivityDefinition) -> None:
     def get_workflow(self, name: str) -> WorkflowDefinition:
         """
         Get a registered workflow by name.
-        
+
         Args:
             name: Name or alias of the workflow
-            
+
         Returns:
-            The workflow definition with type information
-            
+            The workflow definition
+
         Raises:
             KeyError: If workflow is not found
         """
         # Check if it's an alias
         actual_name = self._workflow_aliases.get(name, name)
-        
+
         if actual_name not in self._workflows:
             raise KeyError(f"Workflow '{name}' not found in registry")
-        
+
         return self._workflows[actual_name]
 
     def get_activity(self, name: str) -> ActivityDefinition:

cadence/workflow.py

@@ -2,145 +2,106 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from functools import update_wrapper
-from inspect import signature, Parameter
-from typing import Iterator, Callable, TypeVar, ParamSpec, Generic, TypedDict, Unpack, overload, get_type_hints, Type, Any
+from typing import Iterator, Callable, TypeVar, TypedDict, Type
+from functools import wraps
 
 from cadence.client import Client
 
-
-@dataclass(frozen=True)
-class WorkflowParameter:
-    """Parameter information for a workflow function."""
-    name: str
-    type_hint: Type | None
-    default_value: Any | None
+T = TypeVar('T')
 
 
 class WorkflowDefinitionOptions(TypedDict, total=False):
     """Options for defining a workflow."""
     name: str
 
 
-P = ParamSpec('P')
-T = TypeVar('T')
-
-
-class WorkflowDefinition(Generic[P, T]):
+class WorkflowDefinition:
     """
-    Definition of a workflow function with metadata.
+    Definition of a workflow class with metadata.
 
-    Similar to ActivityDefinition but for workflows.
-    Provides type safety and metadata for workflow functions.
+    Similar to ActivityDefinition but for workflow classes.
+    Provides type safety and metadata for workflow classes.
     """
 
-    def __init__(self, wrapped: Callable[P, T], name: str, params: list[WorkflowParameter]):
-        self._wrapped = wrapped
+    def __init__(self, cls: Type, name: str):
+        self._cls = cls
         self._name = name
-        self._params = params
-        update_wrapper(self, wrapped)
-
-    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
-        return self._wrapped(*args, **kwargs)
 
     @property
     def name(self) -> str:
         """Get the workflow name."""
         return self._name
 
     @property
-    def params(self) -> list[WorkflowParameter]:
-        """Get the workflow parameters."""
-        return self._params
-
-    @property
-    def fn(self) -> Callable[P, T]:
-        """Get the underlying workflow function."""
-        return self._wrapped
+    def cls(self) -> Type:
+        """Get the workflow class."""
+        return self._cls
 
-    @classmethod
-    def wrap(cls, fn: Callable[P, T], opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition[P, T]':
+    @staticmethod
+    def wrap(cls: Type, opts: WorkflowDefinitionOptions) -> 'WorkflowDefinition':
         """
-        Wrap a function as a WorkflowDefinition.
+        Wrap a class as a WorkflowDefinition.
 
         Args:
-            fn: The workflow function to wrap
+            cls: The workflow class to wrap
             opts: Options for the workflow definition
 
         Returns:
             A WorkflowDefinition instance
+
+        Raises:
+            ValueError: If no run method is found or multiple run methods exist
         """
-        name = fn.__qualname__
+        name = cls.__name__
         if "name" in opts and opts["name"]:
             name = opts["name"]
 
-        params = _get_workflow_params(fn)
-        return cls(fn, name, params)
+        # Validate that the class has exactly one run method
+        run_method_count = 0
+        for attr_name in dir(cls):
+            if attr_name.startswith('_'):
+                continue
 
+            attr = getattr(cls, attr_name)
+            if not callable(attr):
+                continue
 
-WorkflowDecorator = Callable[[Callable[P, T]], WorkflowDefinition[P, T]]
+            # Check for workflow run method
+            if hasattr(attr, '_workflow_run'):
+                run_method_count += 1
 
+        if run_method_count == 0:
+            raise ValueError(f"No @workflow.run method found in class {cls.__name__}")
+        elif run_method_count > 1:
+            raise ValueError(f"Multiple @workflow.run methods found in class {cls.__name__}")
 
-@overload
-def defn(fn: Callable[P, T]) -> WorkflowDefinition[P, T]:
-    ...
+        return WorkflowDefinition(cls, name)
 
 
-@overload
-def defn(**kwargs: Unpack[WorkflowDefinitionOptions]) -> WorkflowDecorator:
-    ...
+def run(func: Callable[..., T]) -> Callable[..., T]:
+    """
+    Decorator to mark a method as the main workflow run method.
 
+    Args:
+        func: The method to mark as the workflow run method
 
-def defn(fn: Callable[P, T] | None = None, **kwargs: Unpack[WorkflowDefinitionOptions]) -> WorkflowDecorator | WorkflowDefinition[P, T]:
+    Returns:
+        The decorated method with workflow run metadata
     """
-    Decorator to define a workflow function.
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
 
-    Usage:
-        @defn
-        def my_workflow(input_data: str) -> str:
-            return f"processed: {input_data}"
+    # Attach metadata to the function
+    wrapper._workflow_run = True  # type: ignore
+    return wrapper
 
-        @defn(name="custom_workflow_name")
-        def my_other_workflow(input_data: str) -> str:
-            return f"custom: {input_data}"
 
-    Args:
-        fn: The workflow function (when used without parentheses)
-        **kwargs: Workflow definition options
+# Create a simple namespace object for the workflow decorators
+class _WorkflowNamespace:
+    run = staticmethod(run)
 
-    Returns:
-        Either a WorkflowDefinition (direct decoration) or a decorator function
-    """
-    opts = WorkflowDefinitionOptions(**kwargs)
-
-    def decorator(inner_fn: Callable[P, T]) -> WorkflowDefinition[P, T]:
-        return WorkflowDefinition.wrap(inner_fn, opts)
-
-    if fn is not None:
-        return decorator(fn)
-
-    return decorator
-
-
-def _get_workflow_params(fn: Callable) -> list[WorkflowParameter]:
-    """Extract parameter information from a workflow function."""
-    args = signature(fn).parameters
-    hints = get_type_hints(fn)
-    result = []
-    for name, param in args.items():
-        # Filter out self parameter
-        if param.name == "self":
-            continue
-        default = None
-        if param.default != Parameter.empty:
-            default = param.default
-        if param.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD):
-            type_hint = hints.get(name, None)
-            result.append(WorkflowParameter(name, type_hint, default))
-        else:
-            raise ValueError(f"Parameters must be positional. {name} is {param.kind}, and not valid")
-
-    return result
+workflow = _WorkflowNamespace()
 
 
 @dataclass