Fixing type checking

Author: tconley1428

temporalio/worker/_replayer.py

@@ -7,7 +7,7 @@
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import AsyncIterator, Dict, Mapping, Optional, Sequence, Type, Union
+from typing import AsyncIterator, Dict, Mapping, Optional, Sequence, Type, Union, cast
 
 from typing_extensions import TypedDict
 
@@ -18,12 +18,11 @@
 import temporalio.converter
 import temporalio.runtime
 import temporalio.workflow
-
+from temporalio.client import ClientConfig
 
 from ..common import HeaderCodecBehavior
 from ._interceptor import Interceptor
-from ._worker import load_default_build_id, WorkerConfig
-from temporalio.client import ClientConfig
+from ._worker import WorkerConfig, load_default_build_id
 from ._workflow import _WorkflowWorker
 from ._workflow_instance import UnsandboxedWorkflowRunner, WorkflowRunner
 from .workflow_sandbox import SandboxedWorkflowRunner
@@ -44,7 +43,9 @@ def __init__(
         namespace: str = "ReplayNamespace",
         data_converter: temporalio.converter.DataConverter = temporalio.converter.DataConverter.default,
         interceptors: Sequence[Interceptor] = [],
-        plugins: Sequence[Union[temporalio.worker.Plugin, temporalio.client.Plugin]] = [],
+        plugins: Sequence[
+            Union[temporalio.worker.Plugin, temporalio.client.Plugin]
+        ] = [],
         build_id: Optional[str] = None,
         identity: Optional[str] = None,
         workflow_failure_exception_types: Sequence[Type[BaseException]] = [],
@@ -86,21 +87,59 @@ def __init__(
 
         # Allow plugins to configure shared configurations with worker
         root_worker_plugin: temporalio.worker.Plugin = temporalio.worker._worker._RootPlugin()
-        for plugin in reversed([plugin for plugin in plugins if isinstance(plugin, temporalio.worker.Plugin)]):
+        for plugin in reversed(
+            [
+                plugin
+                for plugin in plugins
+                if isinstance(plugin, temporalio.worker.Plugin)
+            ]
+        ):
             root_worker_plugin = plugin.init_worker_plugin(root_worker_plugin)
 
-        worker_config = WorkerConfig(**{k: v for k, v in self._config.items() if k in WorkerConfig.__annotations__})
+        worker_config = cast(
+            WorkerConfig,
+            {
+                k: v
+                for k, v in self._config.items()
+                if k in WorkerConfig.__annotations__
+            },
+        )
+
         worker_config = root_worker_plugin.configure_worker(worker_config)
-        self._config.update({k: v for k, v in worker_config.items() if k in ReplayerConfig.__annotations__})
+        self._config.update(
+            cast(ReplayerConfig, {
+                k: v
+                for k, v in worker_config.items()
+                if k in ReplayerConfig.__annotations__
+            })
+        )
 
         # Allow plugins to configure shared configurations with client
         root_client_plugin: temporalio.client.Plugin = temporalio.client._RootPlugin()
-        for plugin in reversed([plugin for plugin in plugins if isinstance(plugin, temporalio.client.Plugin)]):
-            root_client_plugin = plugin.init_client_plugin(root_client_plugin)
-
-        client_config = ClientConfig(**{k: v for k, v in self._config.items() if k in ClientConfig.__annotations__})
+        for client_plugin in reversed(
+            [
+                plugin
+                for plugin in plugins
+                if isinstance(plugin, temporalio.client.Plugin)
+            ]
+        ):
+            root_client_plugin = client_plugin.init_client_plugin(root_client_plugin)
+
+        client_config = cast(ClientConfig,
+            {
+                k: v
+                for k, v in self._config.items()
+                if k in ClientConfig.__annotations__
+            }
+        )
         client_config = root_client_plugin.configure_client(client_config)
-        self._config.update({k: v for k, v in client_config.items() if k in ReplayerConfig.__annotations__})
+        self._config.update(
+            cast(ReplayerConfig, {
+                k: v
+                for k, v in client_config.items()
+                if k in ReplayerConfig.__annotations__
+            })
+        )
 
         if not self._config["workflows"]:
             raise ValueError("At least one workflow must be specified")

tests/test_plugins.py

@@ -11,11 +11,10 @@
 from temporalio.client import Client, ClientConfig, OutboundInterceptor
 from temporalio.contrib.pydantic import pydantic_data_converter
 from temporalio.testing import WorkflowEnvironment
-from temporalio.worker import Worker, WorkerConfig
+from temporalio.worker import Replayer, Worker, WorkerConfig
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
-from tests.worker.test_worker import never_run_activity
-from temporalio.worker import Replayer
 from tests.helpers import new_worker
+from tests.worker.test_worker import never_run_activity
 
 
 class TestClientInterceptor(temporalio.client.Interceptor):
@@ -142,21 +141,24 @@ async def test_worker_sandbox_restrictions(client: Client) -> None:
         ).restrictions.passthrough_modules
     )
 
+
 class ReplayCheckPlugin(temporalio.client.Plugin, temporalio.worker.Plugin):
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
-        config["workflows"] = list(config["workflows"]) + [HelloWorkflow]
+        config["workflows"] = list(config.get("workflows") or []) + [HelloWorkflow]
         return super().configure_worker(config)
 
     def configure_client(self, config: ClientConfig) -> ClientConfig:
         config["data_converter"] = pydantic_data_converter
         return super().configure_client(config)
 
+
 @workflow.defn
 class HelloWorkflow:
     @workflow.run
     async def run(self, name: str) -> str:
         return f"Hello, {name}!"
 
+
 async def test_replay(client: Client) -> None:
     plugin = ReplayCheckPlugin()
     new_config = client.config()
@@ -171,28 +173,18 @@ async def test_replay(client: Client) -> None:
             task_queue=worker.task_queue,
         )
         await handle.result()
-    replayer = Replayer(
-        workflows=[],
-        plugins=[plugin]
-    )
-    assert len(replayer.config()["workflows"])==1
-    assert replayer.config()["data_converter"] == pydantic_data_converter
+    replayer = Replayer(workflows=[], plugins=[plugin])
+    assert len(replayer.config().get("workflows") or []) == 1
+    assert replayer.config().get("data_converter") == pydantic_data_converter
 
     await replayer.replay_workflow(await handle.fetch_history())
 
+    replayer = Replayer(workflows=[HelloWorkflow], plugins=[MyClientPlugin()])
+    replayer = Replayer(workflows=[HelloWorkflow], plugins=[MyWorkerPlugin()])
     replayer = Replayer(
-        workflows=[HelloWorkflow],
-        plugins=[MyClientPlugin()]
-    )
-    replayer = Replayer(
-        workflows=[HelloWorkflow],
-        plugins=[MyWorkerPlugin()]
+        workflows=[HelloWorkflow], plugins=[MyClientPlugin(), MyWorkerPlugin()]
     )
     replayer = Replayer(
         workflows=[HelloWorkflow],
-        plugins=[MyClientPlugin(), MyWorkerPlugin()]
+        plugins=[MyWorkerPlugin(), MyClientPlugin(), MyCombinedPlugin()],
     )
-    replayer = Replayer(
-        workflows=[HelloWorkflow],
-        plugins=[MyWorkerPlugin(), MyClientPlugin(), MyCombinedPlugin()]
-    )