Introduce ActivityDefinition as a wrapper for Activity Fns (#29)

Signed-off-by: Nate Mortensen <[email protected]>

Author: natemort

cadence/_internal/activity/_activity_executor.py

@@ -1,4 +1,3 @@
-import inspect
 from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
 from traceback import format_exception
@@ -7,7 +6,7 @@
 from google.protobuf.timestamp import to_datetime
 
 from cadence._internal.activity._context import _Context, _SyncContext
-from cadence.activity import ActivityInfo
+from cadence.activity import ActivityInfo, ActivityDefinition, ExecutionStrategy
 from cadence.api.v1.common_pb2 import Failure
 from cadence.api.v1.service_worker_pb2 import PollForActivityTaskResponse, RespondActivityTaskFailedRequest, \
     RespondActivityTaskCompletedRequest
@@ -16,7 +15,7 @@
 _logger = getLogger(__name__)
 
 class ActivityExecutor:
-    def __init__(self, client: Client, task_list: str, identity: str, max_workers: int, registry: Callable[[str], Callable]):
+    def __init__(self, client: Client, task_list: str, identity: str, max_workers: int, registry: Callable[[str], ActivityDefinition]):
         self._client = client
         self._data_converter = client.data_converter
         self._registry = registry
@@ -36,16 +35,16 @@ async def execute(self, task: PollForActivityTaskResponse):
     def _create_context(self, task: PollForActivityTaskResponse) -> _Context:
         activity_type = task.activity_type.name
         try:
-            activity_fn = self._registry(activity_type)
+            activity_def = self._registry(activity_type)
         except KeyError:
             raise KeyError(f"Activity type not found: {activity_type}") from None
 
         info = self._create_info(task)
 
-        if inspect.iscoroutinefunction(activity_fn):
-            return _Context(self._client, info, activity_fn)
+        if activity_def.strategy == ExecutionStrategy.ASYNC:
+            return _Context(self._client, info, activity_def)
         else:
-            return _SyncContext(self._client, info, activity_fn, self._thread_pool)
+            return _SyncContext(self._client, info, activity_def, self._thread_pool)
 
     async def _report_failure(self, task: PollForActivityTaskResponse, error: Exception):
         try:

cadence/_internal/activity/_context.py

@@ -1,15 +1,14 @@
 import asyncio
 from concurrent.futures.thread import ThreadPoolExecutor
-from typing import Callable, Any
+from typing import Any
 
 from cadence import Client
-from cadence._internal.type_utils import get_fn_parameters
-from cadence.activity import ActivityInfo, ActivityContext
+from cadence.activity import ActivityInfo, ActivityContext, ActivityDefinition
 from cadence.api.v1.common_pb2 import Payload
 
 
 class _Context(ActivityContext):
-    def __init__(self, client: Client, info: ActivityInfo, activity_fn: Callable[[Any], Any]):
+    def __init__(self, client: Client, info: ActivityInfo, activity_fn: ActivityDefinition[[Any], Any]):
         self._client = client
         self._info = info
         self._activity_fn = activity_fn
@@ -20,7 +19,7 @@ async def execute(self, payload: Payload) -> Any:
             return await self._activity_fn(*params)
 
     async def _to_params(self, payload: Payload) -> list[Any]:
-        type_hints = get_fn_parameters(self._activity_fn)
+        type_hints = [param.type_hint for param in self._activity_fn.params]
         return await self._client.data_converter.from_data(payload, type_hints)
 
     def client(self) -> Client:
@@ -30,7 +29,7 @@ def info(self) -> ActivityInfo:
         return self._info
 
 class _SyncContext(_Context):
-    def __init__(self, client: Client, info: ActivityInfo, activity_fn: Callable[[Any], Any], executor: ThreadPoolExecutor):
+    def __init__(self, client: Client, info: ActivityInfo, activity_fn: ActivityDefinition[[Any], Any], executor: ThreadPoolExecutor):
         super().__init__(client, info, activity_fn)
         self._executor = executor
 

cadence/_internal/type_utils.py

@@ -1,9 +1,14 @@
+import inspect
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
 from datetime import timedelta, datetime
-from typing import Iterator
+from enum import Enum
+from functools import update_wrapper
+from inspect import signature, Parameter
+from typing import Iterator, TypedDict, Unpack, Callable, Type, ParamSpec, TypeVar, Generic, get_type_hints, \
+    Any, overload
 
 from cadence import Client
 
@@ -59,3 +64,99 @@ def is_set() -> bool:
     @staticmethod
     def get() -> 'ActivityContext':
         return ActivityContext._var.get()
+
+
+@dataclass(frozen=True)
+class ActivityParameter:
+    name: str
+    type_hint: Type | None
+    default_value: Any | None
+
+class ExecutionStrategy(Enum):
+    ASYNC = "async"
+    THREAD_POOL = "thread_pool"
+
+class ActivityDefinitionOptions(TypedDict, total=False):
+    name: str
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+class ActivityDefinition(Generic[P, T]):
+    def __init__(self, wrapped: Callable[P, T], name: str, strategy: ExecutionStrategy, params: list[ActivityParameter]):
+        self._wrapped = wrapped
+        self._name = name
+        self._strategy = strategy
+        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:
+        return self._name
+
+    @property
+    def strategy(self) -> ExecutionStrategy:
+        return self._strategy
+
+    @property
+    def params(self) -> list[ActivityParameter]:
+        return self._params
+
+    @staticmethod
+    def wrap(fn: Callable[P, T], opts: ActivityDefinitionOptions) -> 'ActivityDefinition[P, T]':
+        name = fn.__qualname__
+        if "name" in opts and opts["name"]:
+            name = opts["name"]
+
+        strategy = ExecutionStrategy.THREAD_POOL
+        if inspect.iscoroutinefunction(fn) or inspect.iscoroutinefunction(fn.__call__): # type: ignore
+            strategy = ExecutionStrategy.ASYNC
+
+        params = _get_params(fn)
+        return ActivityDefinition(fn, name, strategy, params)
+
+
+ActivityDecorator = Callable[[Callable[P, T]], ActivityDefinition[P, T]]
+
+@overload
+def defn(fn: Callable[P, T]) -> ActivityDefinition[P, T]:
+    ...
+
+@overload
+def defn(**kwargs: Unpack[ActivityDefinitionOptions]) -> ActivityDecorator:
+    ...
+
+def defn(fn: Callable[P, T] | None = None, **kwargs: Unpack[ActivityDefinitionOptions]) -> ActivityDecorator | ActivityDefinition[P, T]:
+    options = ActivityDefinitionOptions(**kwargs)
+    def decorator(inner_fn: Callable[P, T]) -> ActivityDefinition[P, T]:
+        return ActivityDefinition.wrap(inner_fn, options)
+
+    if fn is not None:
+        return decorator(fn)
+
+    return decorator
+
+
+def _get_params(fn: Callable) -> list[ActivityParameter]:
+    args = signature(fn).parameters
+    hints = get_type_hints(fn)
+    result = []
+    for name, param in args.items():
+        # "unbound functions" aren't a thing in the Python spec. Filter out the self parameter and hope they followed
+        # the convention.
+        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(ActivityParameter(name, type_hint, default))
+
+        else:
+            raise ValueError(f"Parameters must be positional. {name} is {param.kind}, and not valid")
+
+    return result

cadence/activity.py

@@ -8,13 +8,11 @@
 from ._registry import (
     Registry,
     RegisterWorkflowOptions,
-    RegisterActivityOptions,
 )
 
 __all__ = [
     "Worker",
     "WorkerOptions",
     'Registry',
     'RegisterWorkflowOptions',
-    'RegisterActivityOptions',
 ]