Add optional args= param to ctx.run

python/restate/context.py

@@ -24,7 +24,7 @@
 I = TypeVar('I')
 O = TypeVar('O')
 
-RunAction = Union[Callable[[], T], Callable[[], Awaitable[T]]]
+RunAction = Union[Callable[..., T], Callable[..., Awaitable[T]]]
 
 
 # pylint: disable=R0903
@@ -154,7 +154,8 @@ def run(self,
             serde: Serde[T] = DefaultSerde(),
             max_attempts: typing.Optional[int] = None,
             max_retry_duration: typing.Optional[timedelta] = None,
-            type_hint: Optional[typing.Type[T]] = None
+            type_hint: Optional[typing.Type[T]] = None,
+            args: Optional[typing.Tuple[Any, ...]] = None,
             ) -> RestateDurableFuture[T]:
         """
         Runs the given action with the given name.

python/restate/server_context.py

@@ -19,6 +19,7 @@
 import asyncio
 from datetime import timedelta
 import inspect
+import functools
 from typing import Any, Awaitable, Callable, Dict, List, Optional, TypeVar
 import typing
 import traceback
@@ -397,7 +398,8 @@ async def create_run_coroutine(self,
                                    action: Callable[[], T] | Callable[[], Awaitable[T]],
                                    serde: Serde[T],
                                    max_attempts: Optional[int] = None,
-                                   max_retry_duration: Optional[timedelta] = None):
+                                   max_retry_duration: Optional[timedelta] = None,
+                                   ):
         """Create a coroutine to poll the handle."""
         try:
             if inspect.iscoroutinefunction(action):
@@ -423,20 +425,29 @@ async def create_run_coroutine(self,
     # pylint: disable=R0914
     def run(self,
                   name: str,
-                  action: Callable[[], T] | Callable[[], Awaitable[T]],
+                  action: Callable[..., T] | Callable[..., Awaitable[T]],
                   serde: Serde[T] = DefaultSerde(),
                   max_attempts: Optional[int] = None,
                   max_retry_duration: Optional[timedelta] = None,
-                  type_hint: Optional[typing.Type[T]] = None
+                  type_hint: Optional[typing.Type[T]] = None,
+                  args: Optional[typing.Tuple[Any, ...]] = None
                   ) -> RestateDurableFuture[T]:
 
         if isinstance(serde, DefaultSerde):
             if type_hint is None:
                 signature = inspect.signature(action, eval_str=True)
                 type_hint = signature.return_annotation
             serde = serde.with_maybe_type(type_hint)
+
         handle = self.vm.sys_run(name)
-        self.run_coros_to_execute[handle] = lambda : self.create_run_coroutine(handle, action, serde, max_attempts, max_retry_duration)
+
+        if args is not None:
+            noargs_action = functools.partial(action, *args)
+        else:
+            # todo: we can also verify by looking at the signature that there are no missing parameters
+            noargs_action = action # type: ignore
+
+        self.run_coros_to_execute[handle] = lambda : self.create_run_coroutine(handle, noargs_action, serde, max_attempts, max_retry_duration)
         return self.create_future(handle, serde) # type: ignore