From 952085283952d7ddf2faa959c3e6f7877d6f0e2d Mon Sep 17 00:00:00 2001
From: Dickson Tsai <dickson@anthropic.com>
Date: Sat, 19 Jul 2025 18:47:07 -0700
Subject: [PATCH] Fix types

---
 .../_internal/transport/subprocess_cli.py        |  2 +-
 src/claude_code_sdk/client.py                    | 16 +++++++++-------
 src/claude_code_sdk/query.py                     |  3 ++-
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/claude_code_sdk/_internal/transport/subprocess_cli.py b/src/claude_code_sdk/_internal/transport/subprocess_cli.py
index 92b0743..4121d91 100644
--- a/src/claude_code_sdk/_internal/transport/subprocess_cli.py
+++ b/src/claude_code_sdk/_internal/transport/subprocess_cli.py
@@ -42,7 +42,7 @@ class SubprocessCLITransport(Transport):
         self._stdout_stream: TextReceiveStream | None = None
         self._stderr_stream: TextReceiveStream | None = None
         self._stdin_stream: TextSendStream | None = None
-        self._pending_control_responses: dict[str, Any] = {}
+        self._pending_control_responses: dict[str, dict[str, Any]] = {}
         self._request_counter = 0
         self._close_stdin_after_prompt = close_stdin_after_prompt
 
diff --git a/src/claude_code_sdk/client.py b/src/claude_code_sdk/client.py
index db7c494..213519e 100644
--- a/src/claude_code_sdk/client.py
+++ b/src/claude_code_sdk/client.py
@@ -2,6 +2,7 @@
 
 import os
 from collections.abc import AsyncIterable, AsyncIterator
+from typing import Any
 
 from ._errors import CLIConnectionError
 from .types import ClaudeCodeOptions, Message, ResultMessage
@@ -94,19 +95,20 @@ class ClaudeSDKClient:
         if options is None:
             options = ClaudeCodeOptions()
         self.options = options
-        self._transport = None
+        self._transport: Any | None = None
         os.environ["CLAUDE_CODE_ENTRYPOINT"] = "sdk-py-client"
 
-    async def connect(self, prompt: str | AsyncIterable[dict] | None = None) -> None:
+    async def connect(self, prompt: str | AsyncIterable[dict[str, Any]] | None = None) -> None:
         """Connect to Claude with a prompt or message stream."""
         from ._internal.transport.subprocess_cli import SubprocessCLITransport
 
         # Auto-connect with empty async iterable if no prompt is provided
-        async def _empty_stream():
+        async def _empty_stream() -> AsyncIterator[dict[str, Any]]:
             # Never yields, but indicates that this function is an iterator and
             # keeps the connection open.
-            if False:
-                yield
+            # This yield is never reached but makes this an async generator
+            return
+            yield {}  # type: ignore[unreachable]
 
         self._transport = SubprocessCLITransport(
             prompt=_empty_stream() if prompt is None else prompt,
@@ -190,12 +192,12 @@ class ClaudeSDKClient:
             await self._transport.disconnect()
             self._transport = None
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "ClaudeSDKClient":
         """Enter async context - automatically connects with empty stream for interactive use."""
         await self.connect()
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
         """Exit async context - always disconnects."""
         await self.disconnect()
         return False
diff --git a/src/claude_code_sdk/query.py b/src/claude_code_sdk/query.py
index 4bd7e96..3732762 100644
--- a/src/claude_code_sdk/query.py
+++ b/src/claude_code_sdk/query.py
@@ -2,13 +2,14 @@
 
 import os
 from collections.abc import AsyncIterable, AsyncIterator
+from typing import Any
 
 from ._internal.client import InternalClient
 from .types import ClaudeCodeOptions, Message
 
 
 async def query(
-    *, prompt: str | AsyncIterable[dict], options: ClaudeCodeOptions | None = None
+    *, prompt: str | AsyncIterable[dict[str, Any]], options: ClaudeCodeOptions | None = None
 ) -> AsyncIterator[Message]:
     """
     Query Claude Code for one-shot or unidirectional streaming interactions.