Minor type hint improvments (#95)

* ctx.get should return T instead of Any
* Add additional type hints
* 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

@@ -132,15 +132,15 @@ def get(self,
             name: str,
             serde: Serde[T] = DefaultSerde(),
             type_hint: Optional[typing.Type[T]] = None
-            ) -> Awaitable[Optional[Any]]:
+            ) -> Awaitable[Optional[T]]:
         """
         Retrieves the value associated with the given name.
 
         Args:
             name: The state name
-            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type. 
+            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type.
                     See also 'type_hint'.
-            type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer. 
+            type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer.
         """
 
     @abc.abstractmethod
@@ -213,15 +213,15 @@ def run(self,
         Args:
             name: The name of the action.
             action: The action to run.
-            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type. 
+            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type.
                     See also 'type_hint'.
             max_attempts:   The maximum number of retry attempts to complete the action.
                             If None, the action will be retried indefinitely, until it succeeds.
                             Otherwise, the action will be retried until the maximum number of attempts is reached and then it will raise a TerminalError.
             max_retry_duration: The maximum duration for retrying. If None, the action will be retried indefinitely, until it succeeds.
                                 Otherwise, the action will be retried until the maximum duration is reached and then it will raise a TerminalError.
             type_hint: The type hint of the return value of the action.
-                        This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer. 
+                        This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer.
 
         """
 
@@ -327,7 +327,7 @@ def generic_send(self,
     def awakeable(self,
                   serde: Serde[T] = DefaultSerde(),
                   type_hint: Optional[typing.Type[T]] = None
-                  ) -> typing.Tuple[str, RestateDurableFuture[Any]]:
+                  ) -> typing.Tuple[str, RestateDurableFuture[T]]:
         """
         Returns the name of the awakeable and the future to be awaited.
         """
@@ -388,15 +388,15 @@ def get(self,
             name: str,
             serde: Serde[T] = DefaultSerde(),
             type_hint: Optional[typing.Type[T]] = None
-            ) -> RestateDurableFuture[Optional[Any]]:
+            ) -> RestateDurableFuture[Optional[T]]:
         """
         Retrieves the value associated with the given name.
 
         Args:
             name: The state name
-            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type. 
+            serde: The serialization/deserialization mechanism. - if the default serde is used, a default serializer will be used based on the type.
                     See also 'type_hint'.
-            type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer. 
+            type_hint: The type hint of the return value. This is used to pick the serializer. If None, the type hint will be inferred from the action's return type, or the provided serializer.
         """
 
     @abc.abstractmethod
@@ -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[Any]:
+    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[Any]:
+    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: