type hints for all returns

Author: nagl-temporal

resource_pool/resource_allocator.py

@@ -1,6 +1,6 @@
 from contextlib import asynccontextmanager
 from datetime import timedelta
-from typing import Optional
+from typing import AsyncGenerator, Optional
 
 from temporalio import activity, workflow
 from temporalio.client import Client
@@ -20,7 +20,7 @@
 
 # Use this class in workflow code that that needs to run on locked resources.
 class ResourceAllocator:
-    def __init__(self, client: Client):
+    def __init__(self, client: Client) -> None:
         self.client = client
 
     @activity.defn
@@ -48,13 +48,13 @@ async def acquire_resource(
         *,
         already_acquired_resource: Optional[AcquiredResource] = None,
         max_wait_time: timedelta = timedelta(minutes=5),
-    ):
+    ) -> AsyncGenerator[AcquiredResource, None]:
         warn_when_workflow_has_timeouts()
 
         resource = already_acquired_resource
         if resource is None:
 
-            async def assign_resource(input: AcquireResponse):
+            async def assign_resource(input: AcquireResponse) -> None:
                 workflow.set_signal_handler("assign_resource", None)
                 nonlocal resource
                 resource = AcquiredResource(
@@ -91,7 +91,7 @@ async def assign_resource(input: AcquireResponse):
                 await handle.signal(resource.release_signal_name)
 
 
-def warn_when_workflow_has_timeouts():
+def warn_when_workflow_has_timeouts() -> None:
     if has_timeout(workflow.info().run_timeout):
         workflow.logger.warning(
             f"ResourceLockingWorkflow cannot have a run_timeout (found {workflow.info().run_timeout}) - this will leak locks"

resource_pool/resource_pool_workflow.py

@@ -23,7 +23,7 @@ class ResourcePoolWorkflowInput:
 @workflow.defn
 class ResourcePoolWorkflow:
     @workflow.init
-    def __init__(self, input: ResourcePoolWorkflowInput):
+    def __init__(self, input: ResourcePoolWorkflowInput) -> None:
         self.resources = input.resources
         self.waiters = input.waiters
         self.release_signal_to_resource: dict[str, str] = {}
@@ -32,7 +32,7 @@ def __init__(self, input: ResourcePoolWorkflowInput):
                 self.release_signal_to_resource[holder.release_signal] = resource
 
     @workflow.signal
-    async def add_resources(self, resources: list[str]):
+    async def add_resources(self, resources: list[str]) -> None:
         for resource in resources:
             if resource in self.resources:
                 workflow.logger.warning(
@@ -46,7 +46,7 @@ async def add_resources(self, resources: list[str]):
                 await self.assign_resource(resource, next_holder)
 
     @workflow.signal
-    async def acquire_resource(self, request: AcquireRequest):
+    async def acquire_resource(self, request: AcquireRequest) -> None:
         internal_request = InternalAcquireRequest(
             workflow_id=request.workflow_id, release_signal=None
         )
@@ -65,7 +65,7 @@ async def acquire_resource(self, request: AcquireRequest):
 
     async def assign_resource(
         self, resource: str, internal_request: InternalAcquireRequest
-    ):
+    ) -> None:
         self.resources[resource] = internal_request
         workflow.logger.info(
             f"workflow_id={internal_request.workflow_id} acquired resource {resource}"
@@ -82,7 +82,7 @@ async def assign_resource(
         )
 
     @workflow.signal(dynamic=True)
-    async def release_resource(self, signal_name, *args):
+    async def release_resource(self, signal_name, *args) -> None:
         if not signal_name in self.release_signal_to_resource:
             workflow.logger.warning(
                 f"Ignoring unknown signal: {signal_name} was not a valid release signal."

resource_pool/resource_user_workflow.py

@@ -51,7 +51,7 @@ class FailWorkflowException(Exception):
 @workflow.defn(failure_exception_types=[FailWorkflowException])
 class ResourceUserWorkflow:
     @workflow.run
-    async def run(self, input: ResourceUserWorkflowInput):
+    async def run(self, input: ResourceUserWorkflowInput) -> None:
         async with ResourceAllocator.acquire_resource(
             already_acquired_resource=input.already_acquired_resource
         ) as resource:

resource_pool/starter.py

@@ -15,7 +15,7 @@
 from resource_pool.shared import RESOURCE_POOL_WORKFLOW_ID
 
 
-async def main():
+async def main() -> None:
     # Connect client
     client = await Client.connect("localhost:7233")
 

resource_pool/worker.py

@@ -9,7 +9,7 @@
 from resource_pool.resource_user_workflow import ResourceUserWorkflow, use_resource
 
 
-async def main():
+async def main() -> None:
     logging.basicConfig(level=logging.INFO)
 
     # Start client