From c0c49d31d031a795460ba0a2b27f0dd240d15717 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Fri, 15 Aug 2025 15:57:28 -0700 Subject: [PATCH] Enable Mypy for Workflow (#433) --- .../agent_framework_workflow/_checkpoint.py | 4 ++-- .../agent_framework_workflow/_events.py | 18 +++++++++--------- .../agent_framework_workflow/_executor.py | 2 +- .../agent_framework_workflow/_runner.py | 4 ++-- .../agent_framework_workflow/_shared_state.py | 3 ++- .../agent_framework_workflow/_validation.py | 6 +++--- python/packages/workflow/pyproject.toml | 1 + 7 files changed, 20 insertions(+), 18 deletions(-) diff --git a/python/packages/workflow/agent_framework_workflow/_checkpoint.py b/python/packages/workflow/agent_framework_workflow/_checkpoint.py index 5ef3ae3b06..7814b358dd 100644 --- a/python/packages/workflow/agent_framework_workflow/_checkpoint.py +++ b/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) diff --git a/python/packages/workflow/agent_framework_workflow/_events.py b/python/packages/workflow/agent_framework_workflow/_events.py index 223d7dddd1..857ace18d9 100644 --- a/python/packages/workflow/agent_framework_workflow/_events.py +++ b/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})" diff --git a/python/packages/workflow/agent_framework_workflow/_executor.py b/python/packages/workflow/agent_framework_workflow/_executor.py index b1ddad61b6..cac7e73f6e 100644 --- a/python/packages/workflow/agent_framework_workflow/_executor.py +++ b/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] = {} diff --git a/python/packages/workflow/agent_framework_workflow/_runner.py b/python/packages/workflow/agent_framework_workflow/_runner.py index ac6ca41fb1..4235e6842d 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner.py +++ b/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.""" diff --git a/python/packages/workflow/agent_framework_workflow/_shared_state.py b/python/packages/workflow/agent_framework_workflow/_shared_state.py index a43bcfc30b..22225518ea 100644 --- a/python/packages/workflow/agent_framework_workflow/_shared_state.py +++ b/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: diff --git a/python/packages/workflow/agent_framework_workflow/_validation.py b/python/packages/workflow/agent_framework_workflow/_validation.py index 3c4d8c12fd..581475a396 100644 --- a/python/packages/workflow/agent_framework_workflow/_validation.py +++ b/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 diff --git a/python/packages/workflow/pyproject.toml b/python/packages/workflow/pyproject.toml index 73577d3a45..d9358c80e3 100644 --- a/python/packages/workflow/pyproject.toml +++ b/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]