more type hints around durable promise

Author: igalshilman

examples/virtual_object.py

@@ -18,7 +18,7 @@
 
 @counter.handler()
 async def increment(ctx: ObjectContext, value: int) -> int:
-    n = await ctx.get("counter") or 0
+    n = await ctx.get("counter", type_hint=int) or 0
     n += value
     ctx.set("counter", n)
     return n

examples/workflow.py

@@ -56,5 +56,5 @@ def payment_gateway():
 
 @payment.handler()
 async def payment_verified(ctx: WorkflowSharedContext, result: str):
-    promise = ctx.promise("verify.payment")
+    promise = ctx.promise("verify.payment", type_hint=str)
     await promise.resolve(result)

python/restate/context.py

@@ -438,13 +438,19 @@ def value(self) -> RestateDurableFuture[T]:
         Returns the value of the promise if it is resolved, None otherwise.
         """
 
+    @abc.abstractmethod
+    def __await__(self) -> typing.Generator[Any, Any, T]:
+        """
+        Returns the value of the promise. This is a shortcut for calling value() and awaiting it. 
+        """
+
 class WorkflowContext(ObjectContext):
     """
     Represents the context of the current workflow invocation.
     """
 
     @abc.abstractmethod
-    def promise(self, name: str, serde: Serde[T] = DefaultSerde()) -> DurablePromise[T]:
+    def promise(self, name: str, serde: Serde[T] = DefaultSerde(), type_hint: Optional[typing.Type[T]] = None) -> DurablePromise[T]:
         """
         Returns a durable promise with the given name.
         """
@@ -455,7 +461,7 @@ class WorkflowSharedContext(ObjectSharedContext):
     """
 
     @abc.abstractmethod
-    def promise(self, name: str, serde: Serde[T] = DefaultSerde()) -> DurablePromise[T]:
+    def promise(self, name: str, serde: Serde[T] = DefaultSerde(), type_hint: Optional[typing.Type[T]] = None) -> DurablePromise[T]:
         """
         Returns a durable promise with the given name.
         """

python/restate/server_context.py

@@ -216,6 +216,9 @@ def peek(self) -> Awaitable[Any | None]:
         assert serde is not None
         return self.server_context.create_future(handle, serde)
 
+    def __await__(self):
+        return self.value().__await__()
+
 
 # disable too many public method
 # pylint: disable=R0904
@@ -684,8 +687,10 @@ def resolve_awakeable(self,
     def reject_awakeable(self, name: str, failure_message: str, failure_code: int = 500) -> None:
         return self.vm.sys_reject_awakeable(name, Failure(code=failure_code, message=failure_message))
 
-    def promise(self, name: str, serde: typing.Optional[Serde[T]] = JsonSerde()) -> DurablePromise[Any]:
+    def promise(self, name: str, serde: typing.Optional[Serde[T]] = JsonSerde(), type_hint: Optional[typing.Type[T]] = None) -> DurablePromise[T]:
         """Create a durable promise."""
+        if isinstance(serde, DefaultSerde):
+            serde = serde.with_maybe_type(type_hint)
         return ServerDurablePromise(self, name, serde)
 
     def key(self) -> str: