Improve entity typing

Author: andystaples

durabletask/client.py

@@ -11,8 +11,7 @@
 import grpc
 from google.protobuf import wrappers_pb2
 
-from durabletask.entities import EntityInstanceId
-from durabletask.entities.entity_metadata import EntityMetadata
+from durabletask.entities import EntityInstanceId, EntityMetadata
 import durabletask.internal.helpers as helpers
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
@@ -230,7 +229,10 @@ def purge_orchestration(self, instance_id: str, recursive: bool = True):
         self._logger.info(f"Purging instance '{instance_id}'.")
         self._stub.PurgeInstances(req)
 
-    def signal_entity(self, entity_instance_id: EntityInstanceId, operation_name: str, input: Optional[Any] = None):
+    def signal_entity(self,
+                      entity_instance_id: EntityInstanceId[TInput, TOutput],
+                      operation_name: str,
+                      input: Optional[TInput] = None):
         req = pb.SignalEntityRequest(
             instanceId=str(entity_instance_id),
             name=operation_name,
@@ -244,7 +246,7 @@ def signal_entity(self, entity_instance_id: EntityInstanceId, operation_name: st
         self._stub.SignalEntity(req, None)  # TODO: Cancellation timeout?
 
     def get_entity(self,
-                   entity_instance_id: EntityInstanceId,
+                   entity_instance_id: EntityInstanceId[Any, Any],
                    include_state: bool = True
                    ) -> Optional[EntityMetadata]:
         req = pb.GetEntityRequest(instanceId=str(entity_instance_id), includeState=include_state)

durabletask/entities/durable_entity.py

@@ -4,6 +4,7 @@
 from durabletask.entities.entity_instance_id import EntityInstanceId
 
 TState = TypeVar("TState")
+TInput = TypeVar("TInput")
 
 
 class DurableEntity:
@@ -49,7 +50,10 @@ def set_state(self, state: Any):
         """
         self.entity_context.set_state(state)
 
-    def signal_entity(self, entity_instance_id: EntityInstanceId, operation: str, input: Optional[Any] = None) -> None:
+    def signal_entity(self,
+                      entity_instance_id: EntityInstanceId[TInput, Any],
+                      operation: str,
+                      input: Optional[TInput] = None) -> None:
         """Signal another entity to perform an operation.
 
         Parameters

durabletask/entities/entity_context.py

@@ -7,6 +7,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 
 TState = TypeVar("TState")
+TInput = TypeVar("TInput")
 
 
 class EntityContext:
@@ -81,7 +82,7 @@ def set_state(self, new_state: Any):
         """
         self._state.set_state(new_state)
 
-    def signal_entity(self, entity_instance_id: EntityInstanceId, operation: str, input: Optional[Any] = None) -> None:
+    def signal_entity(self, entity_instance_id: EntityInstanceId[TInput, Any], operation: str, input: Optional[TInput] = None) -> None:
         """Signal another entity to perform an operation.
 
         Parameters

durabletask/entities/entity_instance_id.py

@@ -1,7 +1,54 @@
-class EntityInstanceId:
-    def __init__(self, entity: str, key: str):
-        self.entity = entity
-        self.key = key
+from typing import Any, Callable, TypeVar, Union, overload, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from durabletask import task
+    from durabletask.entities.durable_entity import DurableEntity
+    from durabletask.entities.entity_context import EntityContext
+
+
+TInput = TypeVar('TInput')
+TOutput = TypeVar('TOutput')
+
+
+class EntityInstanceId[TInput, TOutput]:
+    @overload
+    def __new__(
+        cls,
+        entity: Callable[[EntityContext, TInput], TOutput],
+        key: str
+    ) -> "EntityInstanceId[TInput, TOutput]": ...
+
+    @overload
+    def __new__(
+        cls,
+        entity: type[DurableEntity],
+        key: str
+    ) -> "EntityInstanceId[Any, Any]": ...
+
+    @overload
+    def __new__(
+        cls,
+        entity: str,
+        key: str
+    ) -> "EntityInstanceId[Any, Any]": ...
+
+    def __new__(
+        cls,
+        entity: Union[task.Entity[TInput, TOutput], str],
+        key: str
+    ) -> "EntityInstanceId[Any, Any]":
+        return super().__new__(cls)
+
+    def __init__(
+        self,
+        entity: Union[task.Entity[TInput, TOutput], str],
+        key: str
+    ):
+        if not isinstance(entity, str):
+            from durabletask import task
+            entity = task.get_entity_name(entity)
+        self.entity: str = entity
+        self.key: str = key
 
     def __str__(self) -> str:
         return f"@{self.entity}@{self.key}"
@@ -17,7 +64,7 @@ def __lt__(self, other):
         return str(self) < str(other)
 
     @staticmethod
-    def parse(entity_id: str) -> "EntityInstanceId":
+    def parse(entity_id: str) -> "EntityInstanceId[Any, Any]":
         """Parse a string representation of an entity ID into an EntityInstanceId object.
 
         Parameters

durabletask/entities/entity_metadata.py

@@ -24,7 +24,7 @@ class EntityMetadata:
     """
 
     def __init__(self,
-                 id: EntityInstanceId,
+                 id: EntityInstanceId[Any, Any],
                  last_modified: datetime,
                  backlog_queue_size: int,
                  locked_by: str,

durabletask/task.py

@@ -140,14 +140,14 @@ def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
 
     @abstractmethod
     def call_entity(self,
-                    entity: EntityInstanceId,
+                    entity: EntityInstanceId[TInput, TOutput],
                     operation: str,
-                    input: Optional[TInput] = None) -> CompletableTask:
+                    input: Optional[TInput] = None) -> CompletableTask[TOutput]:
         """Schedule entity function for execution.
 
         Parameters
         ----------
-        entity: EntityInstanceId
+        entity: EntityInstanceId[TInput, TOutput]
             The ID of the entity instance to call.
         operation: str
             The name of the operation to invoke on the entity.
@@ -164,15 +164,15 @@ def call_entity(self,
     @abstractmethod
     def signal_entity(
             self,
-            entity_id: EntityInstanceId,
+            entity_id: EntityInstanceId[TInput, TOutput],
             operation_name: str,
             input: Optional[TInput] = None
     ) -> None:
         """Signal an entity function for execution.
 
         Parameters
         ----------
-        entity_id: EntityInstanceId
+        entity_id: EntityInstanceId[TInput, TOutput]
             The ID of the entity instance to signal.
         operation_name: str
             The name of the operation to invoke on the entity.
@@ -182,7 +182,7 @@ def signal_entity(
         pass
 
     @abstractmethod
-    def lock_entities(self, entities: list[EntityInstanceId]) -> CompletableTask[EntityLock]:
+    def lock_entities(self, entities: list[EntityInstanceId[Any, Any]]) -> CompletableTask[EntityLock]:
         """Creates a Task object that locks the specified entity instances.
 
         The locks will be acquired the next time the orchestrator yields.
@@ -191,7 +191,7 @@ def lock_entities(self, entities: list[EntityInstanceId]) -> CompletableTask[Ent
 
         Parameters
         ----------
-        entities: list[EntityInstanceId]
+        entities: list[EntityInstanceId[Any, Any]]
             The list of entity instance IDs to lock.
 
         Returns
@@ -538,8 +538,8 @@ def task_id(self) -> int:
         return self._task_id
 
 
-# Orchestrators are generators that yield tasks and receive/return any type
-Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, Any, Any], TOutput]]
+# Orchestrators are generators that yield tasks, recieve any type, and return TOutput
+Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task[Any], Any, TOutput], TOutput]]
 
 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]
@@ -615,6 +615,14 @@ def retry_timeout(self) -> Optional[timedelta]:
         return self._retry_timeout
 
 
+def get_entity_name(fn: Entity) -> str:
+    if hasattr(fn, "__durable_entity_name__"):
+        return getattr(fn, "__durable_entity_name__")
+    if isinstance(fn, type) and issubclass(fn, DurableEntity):
+        return fn.__name__
+    return get_name(fn)
+
+
 def get_name(fn: Callable) -> str:
     """Returns the name of the provided function"""
     name = fn.__name__

durabletask/worker.py

@@ -150,23 +150,23 @@ def __init__(self):
         self.entities = {}
         self.entity_instances = {}
 
-    def add_orchestrator(self, fn: task.Orchestrator) -> str:
+    def add_orchestrator(self, fn: task.Orchestrator[TInput, TOutput]) -> str:
         if fn is None:
             raise ValueError("An orchestrator function argument is required.")
 
         name = task.get_name(fn)
         self.add_named_orchestrator(name, fn)
         return name
 
-    def add_named_orchestrator(self, name: str, fn: task.Orchestrator) -> None:
+    def add_named_orchestrator(self, name: str, fn: task.Orchestrator[TInput, TOutput]) -> None:
         if not name:
             raise ValueError("A non-empty orchestrator name is required.")
         if name in self.orchestrators:
             raise ValueError(f"A '{name}' orchestrator already exists.")
 
         self.orchestrators[name] = fn
 
-    def get_orchestrator(self, name: str) -> Optional[task.Orchestrator]:
+    def get_orchestrator(self, name: str) -> Optional[task.Orchestrator[Any, Any]]:
         return self.orchestrators.get(name)
 
     def add_activity(self, fn: task.Activity) -> str:
@@ -188,16 +188,13 @@ def add_named_activity(self, name: str, fn: task.Activity) -> None:
     def get_activity(self, name: str) -> Optional[task.Activity]:
         return self.activities.get(name)
 
-    def add_entity(self, fn: task.Entity) -> str:
+    def add_entity(self, fn: task.Entity, name: Optional[str] = None) -> str:
         if fn is None:
             raise ValueError("An entity function argument is required.")
 
-        if isinstance(fn, type) and issubclass(fn, DurableEntity):
-            name = fn.__name__
-            self.add_named_entity(name, fn)
-        else:
-            name = task.get_name(fn)
-            self.add_named_entity(name, fn)
+        if name is None:
+            name = task.get_entity_name(fn)
+        self.add_named_entity(name, fn)
         return name
 
     def add_named_entity(self, name: str, fn: task.Entity) -> None:
@@ -207,6 +204,7 @@ def add_named_entity(self, name: str, fn: task.Entity) -> None:
             raise ValueError(f"A '{name}' entity already exists.")
 
         self.entities[name] = fn
+        setattr(fn, "__durable_entity_name__", name)
 
     def get_entity(self, name: str) -> Optional[task.Entity]:
         return self.entities.get(name)
@@ -362,7 +360,7 @@ def __enter__(self):
     def __exit__(self, type, value, traceback):
         self.stop()
 
-    def add_orchestrator(self, fn: task.Orchestrator) -> str:
+    def add_orchestrator(self, fn: task.Orchestrator[TInput, TOutput]) -> str:
         """Registers an orchestrator function with the worker."""
         if self._is_running:
             raise RuntimeError(
@@ -378,13 +376,13 @@ def add_activity(self, fn: task.Activity) -> str:
             )
         return self._registry.add_activity(fn)
 
-    def add_entity(self, fn: task.Entity) -> str:
+    def add_entity(self, fn: task.Entity, name: Optional[str] = None) -> str:
         """Registers an entity function with the worker."""
         if self._is_running:
             raise RuntimeError(
                 "Entities cannot be added while the worker is running."
             )
-        return self._registry.add_entity(fn)
+        return self._registry.add_entity(fn, name)
 
     def use_versioning(self, version: VersioningOptions) -> None:
         """Initializes versioning options for sub-orchestrators and activities."""
@@ -1044,21 +1042,21 @@ def call_activity(
 
     def call_entity(
             self,
-            entity: EntityInstanceId,
+            entity: EntityInstanceId[TInput, TOutput],
             operation: str,
             input: Optional[TInput] = None,
-    ) -> task.CompletableTask:
+    ) -> task.CompletableTask[TOutput]:
         id = self.next_sequence_number()
 
         self.call_entity_function_helper(
             id, entity, operation, input=input
         )
 
-        return self._pending_tasks.get(id, task.CompletableTask())
+        return self._pending_tasks.get(id, task.CompletableTask[TOutput]())
 
     def signal_entity(
             self,
-            entity_id: EntityInstanceId,
+            entity_id: EntityInstanceId[TInput, TOutput],
             operation_name: str,
             input: Optional[TInput] = None
     ) -> None:
@@ -1068,7 +1066,7 @@ def signal_entity(
             id, entity_id, operation_name, input
         )
 
-    def lock_entities(self, entities: list[EntityInstanceId]) -> task.CompletableTask[EntityLock]:
+    def lock_entities(self, entities: list[EntityInstanceId[Any, Any]]) -> task.CompletableTask[EntityLock]:
         id = self.next_sequence_number()
 
         self.lock_entities_function_helper(
@@ -1158,11 +1156,11 @@ def call_activity_function_helper(
     def call_entity_function_helper(
             self,
             id: Optional[int],
-            entity_id: EntityInstanceId,
+            entity_id: EntityInstanceId[TInput, TOutput],
             operation: str,
             *,
             input: Optional[TInput] = None,
-    ):
+    ) -> None:
         if id is None:
             id = self.next_sequence_number()
 
@@ -1180,7 +1178,7 @@ def call_entity_function_helper(
     def signal_entity_function_helper(
             self,
             id: Optional[int],
-            entity_id: EntityInstanceId,
+            entity_id: EntityInstanceId[TInput, TOutput],
             operation: str,
             input: Optional[TInput]
     ) -> None:
@@ -1197,7 +1195,7 @@ def signal_entity_function_helper(
         action = ph.new_signal_entity_action(id, entity_id, operation, encoded_input, self.new_uuid())
         self._pending_actions[id] = action
 
-    def lock_entities_function_helper(self, id: int, entities: list[EntityInstanceId]) -> None:
+    def lock_entities_function_helper(self, id: int, entities: list[EntityInstanceId[Any, Any]]) -> None:
         if id is None:
             id = self.next_sequence_number()
 
@@ -1792,7 +1790,7 @@ def process_event(
             # The orchestrator generator function completed
             ctx.set_complete(generatorStopped.value, pb.ORCHESTRATION_STATUS_COMPLETED)
 
-    def _parse_entity_event_sent_input(self, event: pb.HistoryEvent) -> Tuple[EntityInstanceId, str]:
+    def _parse_entity_event_sent_input(self, event: pb.HistoryEvent) -> Tuple[EntityInstanceId[Any, Any], str]:
         try:
             entity_id = EntityInstanceId.parse(event.eventSent.instanceId)
         except ValueError:
@@ -1806,7 +1804,7 @@ def _parse_entity_event_sent_input(self, event: pb.HistoryEvent) -> Tuple[Entity
     def _handle_entity_event_raised(self,
                                     ctx: _RuntimeOrchestrationContext,
                                     event: pb.HistoryEvent,
-                                    entity_id: Optional[EntityInstanceId],
+                                    entity_id: Optional[EntityInstanceId[Any, Any]],
                                     task_id: Optional[int],
                                     is_lock_event: bool):
         # This eventRaised represents the result of an entity operation after being translated to the old
@@ -1919,7 +1917,7 @@ def __init__(self, registry: _Registry, logger: logging.Logger):
     def execute(
             self,
             orchestration_id: str,
-            entity_id: EntityInstanceId,
+            entity_id: EntityInstanceId[TInput, TOutput],
             operation: str,
             state: StateShim,
             encoded_input: Optional[str],