Cleanup

Author: dandavison

temporalio/worker/_command_aware_visitor.py

@@ -3,11 +3,10 @@
 import contextvars
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Iterator, Optional, Type
+from typing import Iterator, Optional
 
 from temporalio.api.enums.v1.command_type_pb2 import CommandType
 from temporalio.bridge._visitor import PayloadVisitor, VisitorFunctions
-from temporalio.bridge.proto.workflow_activation import workflow_activation_pb2
 from temporalio.bridge.proto.workflow_activation.workflow_activation_pb2 import (
     ResolveActivity,
     ResolveChildWorkflowExecution,
@@ -17,7 +16,6 @@
     ResolveRequestCancelExternalWorkflow,
     ResolveSignalExternalWorkflow,
 )
-from temporalio.bridge.proto.workflow_commands import workflow_commands_pb2
 from temporalio.bridge.proto.workflow_commands.workflow_commands_pb2 import (
     ScheduleActivity,
     ScheduleLocalActivity,
@@ -150,27 +148,13 @@ async def _visit_coresdk_workflow_activation_ResolveNexusOperation(
             )
 
 
-def _get_workflow_command_protos_with_seq() -> Iterator[Type[Any]]:
-    """Get concrete classes of all workflow command protos with a seq field."""
-    for descriptor in workflow_commands_pb2.DESCRIPTOR.message_types_by_name.values():
-        if "seq" in descriptor.fields_by_name:
-            yield descriptor._concrete_class
-
-
-def _get_workflow_activation_job_protos_with_seq() -> Iterator[Type[Any]]:
-    """Get concrete classes of all workflow activation job protos with a seq field."""
-    for descriptor in workflow_activation_pb2.DESCRIPTOR.message_types_by_name.values():
-        if "seq" in descriptor.fields_by_name:
-            yield descriptor._concrete_class
-
-
 @contextmanager
 def current_command(
-    command_type: Optional[CommandType.ValueType], command_seq: int
+    command_type: CommandType.ValueType, command_seq: int
 ) -> Iterator[None]:
     """Context manager for setting command info."""
     token = current_command_info.set(
-        CommandInfo(command_type=command_type, command_seq=command_seq)  # type: ignore
+        CommandInfo(command_type=command_type, command_seq=command_seq)
     )
     try:
         yield

tests/worker/test_command_aware_visitor.py

@@ -1,11 +1,11 @@
 """Test that CommandAwarePayloadVisitor handles all commands with seq fields that have payloads."""
 
+from typing import Any, Iterator, Type
+
 from temporalio.bridge._visitor import PayloadVisitor
-from temporalio.worker._command_aware_visitor import (
-    CommandAwarePayloadVisitor,
-    _get_workflow_activation_job_protos_with_seq,
-    _get_workflow_command_protos_with_seq,
-)
+from temporalio.bridge.proto.workflow_activation import workflow_activation_pb2
+from temporalio.bridge.proto.workflow_commands import workflow_commands_pb2
+from temporalio.worker._command_aware_visitor import CommandAwarePayloadVisitor
 
 
 def test_command_aware_visitor_has_methods_for_all_seq_protos_with_payloads():
@@ -73,3 +73,17 @@ def test_command_aware_visitor_has_methods_for_all_seq_protos_with_payloads():
     assert (
         len(jobs_with_payloads) == len(job_protos) - 1
     ), "Should have exactly one activation job without payloads (FireTimer)"
+
+
+def _get_workflow_command_protos_with_seq() -> Iterator[Type[Any]]:
+    """Get concrete classes of all workflow command protos with a seq field."""
+    for descriptor in workflow_commands_pb2.DESCRIPTOR.message_types_by_name.values():
+        if "seq" in descriptor.fields_by_name:
+            yield descriptor._concrete_class
+
+
+def _get_workflow_activation_job_protos_with_seq() -> Iterator[Type[Any]]:
+    """Get concrete classes of all workflow activation job protos with a seq field."""
+    for descriptor in workflow_activation_pb2.DESCRIPTOR.message_types_by_name.values():
+        if "seq" in descriptor.fields_by_name:
+            yield descriptor._concrete_class