Enable Mypy for Workflow (#433)

python/packages/workflow/agent_framework_workflow/_checkpoint.py

@@ -71,7 +71,7 @@ class CheckpointStorage(Protocol):
 class InMemoryCheckpointStorage:
     """In-memory checkpoint storage for testing and development."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the memory storage."""
         self._checkpoints: dict[str, WorkflowCheckpoint] = {}
 
@@ -143,7 +143,7 @@ class FileCheckpointStorage:
 
         def _read() -> dict[str, Any]:
             with open(file_path) as f:
-                return json.load(f)
+                return json.load(f)  # type: ignore[no-any-return]
 
         checkpoint_dict = await asyncio.to_thread(_read)
 

python/packages/workflow/agent_framework_workflow/_events.py

@@ -12,7 +12,7 @@ class WorkflowEvent:
         """Initialize the workflow event with optional data."""
         self.data = data
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the workflow event."""
         return f"{self.__class__.__name__}(data={self.data if self.data is not None else 'None'})"
 
@@ -36,7 +36,7 @@ class WorkflowWarningEvent(WorkflowEvent):
         """Initialize the workflow warning event with optional data and warning message."""
         super().__init__(data)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the workflow warning event."""
         return f"{self.__class__.__name__}(message={self.data})"
 
@@ -48,7 +48,7 @@ class WorkflowErrorEvent(WorkflowEvent):
         """Initialize the workflow error event with optional data and error message."""
         super().__init__(data)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the workflow error event."""
         return f"{self.__class__.__name__}(exception={self.data})"
 
@@ -76,7 +76,7 @@ class RequestInfoEvent(WorkflowEvent):
         self.source_executor_id = source_executor_id
         self.request_type = request_type
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the request info event."""
         return (
             f"{self.__class__.__name__}("
@@ -95,7 +95,7 @@ class ExecutorEvent(WorkflowEvent):
         super().__init__(data)
         self.executor_id = executor_id
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the executor event."""
         return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"
 
@@ -103,7 +103,7 @@ class ExecutorEvent(WorkflowEvent):
 class ExecutorInvokeEvent(ExecutorEvent):
     """Event triggered when an executor handler is invoked."""
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the executor handler invoke event."""
         return f"{self.__class__.__name__}(executor_id={self.executor_id})"
 
@@ -111,7 +111,7 @@ class ExecutorInvokeEvent(ExecutorEvent):
 class ExecutorCompletedEvent(ExecutorEvent):
     """Event triggered when an executor handler is completed."""
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the executor handler complete event."""
         return f"{self.__class__.__name__}(executor_id={self.executor_id})"
 
@@ -123,7 +123,7 @@ class AgentRunStreamingEvent(ExecutorEvent):
         """Initialize the agent streaming event."""
         super().__init__(executor_id, data)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the agent streaming event."""
         return f"{self.__class__.__name__}(executor_id={self.executor_id}, messages={self.data})"
 
@@ -135,6 +135,6 @@ class AgentRunEvent(ExecutorEvent):
         """Initialize the agent run event."""
         super().__init__(executor_id, data)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Return a string representation of the agent run event."""
         return f"{self.__class__.__name__}(executor_id={self.executor_id}, data={self.data})"

python/packages/workflow/agent_framework_workflow/_executor.py

@@ -294,7 +294,7 @@ class RequestInfoExecutor(Executor):
     # Well-known ID for the request info executor
     EXECUTOR_ID: ClassVar[str] = "request_info"
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize the RequestInfoExecutor with its well-known ID."""
         super().__init__(id=self.EXECUTOR_ID)
         self._request_events: dict[str, RequestInfoEvent] = {}

python/packages/workflow/agent_framework_workflow/_runner.py

@@ -25,7 +25,7 @@ class Runner:
         ctx: RunnerContext,
         max_iterations: int = 100,
         workflow_id: str | None = None,
-    ):
+    ) -> None:
         """Initialize the runner with edges, shared state, and context.
 
         Args:
@@ -123,7 +123,7 @@ class Runner:
         finally:
             self._running = False
 
-    async def _run_iteration(self):
+    async def _run_iteration(self) -> None:
         async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None:
             """Outer loop to concurrently deliver messages from all sources to their targets."""
 

python/packages/workflow/agent_framework_workflow/_shared_state.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from typing import Any
 
@@ -34,7 +35,7 @@ class SharedState:
             await self.delete_within_hold(key)
 
     @asynccontextmanager
-    async def hold(self):
+    async def hold(self) -> AsyncIterator["SharedState"]:
         """Context manager to hold the shared state lock for multiple operations.
 
         Usage:

python/packages/workflow/agent_framework_workflow/_validation.py

@@ -88,7 +88,7 @@ class WorkflowGraphValidator:
     3. Graph connectivity validation
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._edges: list[Edge] = []
         self._executors: dict[str, Executor] = {}
 
@@ -209,9 +209,9 @@ class WorkflowGraphValidator:
             for target_type in target_input_types:
                 if isinstance(edge_group, FanInEdgeGroup):
                     # If the edge is part of an edge group, the target expects a list of data types
-                    if self._is_type_compatible(list[source_type], target_type):
+                    if self._is_type_compatible(list[source_type], target_type):  # type: ignore[valid-type]
                         compatible = True
-                        compatible_pairs.append((list[source_type], target_type))
+                        compatible_pairs.append((list[source_type], target_type))  # type: ignore[valid-type]
                 else:
                     if self._is_type_compatible(source_type, target_type):
                         compatible = True

python/packages/workflow/pyproject.toml

@@ -79,6 +79,7 @@ exclude_dirs = ["tests"]
 executor.type = "uv"
 include = "../../shared_tasks.toml"
 [tool.poe.tasks]
+mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_workflow"
 test = "pytest --cov=agent_framework_workflow --cov-report=term-missing:skip-covered tests"
 
 [tool.uv.build-backend]