Fix code formatting for CI

This commit is contained in:
Lina Tawfik 2025-06-12 00:20:28 -07:00
parent 6ca3514261
commit 63ef121e18
No known key found for this signature in database
8 changed files with 101 additions and 79 deletions

View file

@ -34,8 +34,7 @@ async def with_options_example():
)
async for message in query(
prompt="Explain what Python is in one sentence.",
options=options
prompt="Explain what Python is in one sentence.", options=options
):
if isinstance(message, AssistantMessage):
for block in message.content:
@ -55,7 +54,7 @@ async def with_tools_example():
async for message in query(
prompt="Create a file called hello.txt with 'Hello, World!' in it",
options=options
options=options,
):
if isinstance(message, AssistantMessage):
for block in message.content:

View file

@ -48,9 +48,7 @@ class InternalClient:
match data["type"]:
case "user":
# Extract just the content from the nested structure
return UserMessage(
content=data["message"]["content"]
)
return UserMessage(content=data["message"]["content"])
case "assistant":
# Parse content blocks
@ -60,24 +58,28 @@ class InternalClient:
case "text":
content_blocks.append(TextBlock(text=block["text"]))
case "tool_use":
content_blocks.append(ToolUseBlock(
id=block["id"],
name=block["name"],
input=block["input"]
))
content_blocks.append(
ToolUseBlock(
id=block["id"],
name=block["name"],
input=block["input"],
)
)
case "tool_result":
content_blocks.append(ToolResultBlock(
tool_use_id=block["tool_use_id"],
content=block.get("content"),
is_error=block.get("is_error")
))
content_blocks.append(
ToolResultBlock(
tool_use_id=block["tool_use_id"],
content=block.get("content"),
is_error=block.get("is_error"),
)
)
return AssistantMessage(content=content_blocks)
case "system":
return SystemMessage(
subtype=data["subtype"],
data=data # Pass through all data
data=data, # Pass through all data
)
case "result":
@ -92,7 +94,7 @@ class InternalClient:
session_id=data["session_id"],
total_cost_usd=data["total_cost"],
usage=data.get("usage"),
result=data.get("result")
result=data.get("result"),
)
case _:

View file

@ -22,7 +22,10 @@ class SubprocessCLITransport(Transport):
"""Subprocess transport using Claude Code CLI."""
def __init__(
self, prompt: str, options: ClaudeCodeOptions, cli_path: str | Path | None = None
self,
prompt: str,
options: ClaudeCodeOptions,
cli_path: str | Path | None = None,
):
self._prompt = prompt
self._options = options

View file

@ -13,6 +13,7 @@ PermissionMode = Literal["default", "acceptEdits", "bypassPermissions"]
# MCP Server config
class McpServerConfig(TypedDict):
"""MCP server configuration."""
transport: list[str]
env: NotRequired[dict[str, Any]]
@ -21,12 +22,14 @@ class McpServerConfig(TypedDict):
@dataclass
class TextBlock:
"""Text content block."""
text: str
@dataclass
class ToolUseBlock:
"""Tool use content block."""
id: str
name: str
input: dict[str, Any]
@ -35,6 +38,7 @@ class ToolUseBlock:
@dataclass
class ToolResultBlock:
"""Tool result content block."""
tool_use_id: str
content: str | list[dict[str, Any]] | None = None
is_error: bool | None = None
@ -47,18 +51,21 @@ ContentBlock = TextBlock | ToolUseBlock | ToolResultBlock
@dataclass
class UserMessage:
"""User message."""
content: str
@dataclass
class AssistantMessage:
"""Assistant message with content blocks."""
content: list[ContentBlock]
@dataclass
class SystemMessage:
"""System message with metadata."""
subtype: str
data: dict[str, Any]
@ -66,6 +73,7 @@ class SystemMessage:
@dataclass
class ResultMessage:
"""Result message with cost and usage information."""
subtype: str
cost_usd: float
duration_ms: int
@ -84,6 +92,7 @@ Message = UserMessage | AssistantMessage | SystemMessage | ResultMessage
@dataclass
class ClaudeCodeOptions:
"""Query options for Claude SDK."""
allowed_tools: list[str] = field(default_factory=list)
max_thinking_tokens: int = 8000
system_prompt: str | None = None

View file

@ -13,13 +13,14 @@ class TestQueryFunction:
def test_query_single_prompt(self):
"""Test query with a single prompt."""
async def _test():
with patch('claude_code_sdk._internal.client.InternalClient.process_query') as mock_process:
with patch(
"claude_code_sdk._internal.client.InternalClient.process_query"
) as mock_process:
# Mock the async generator
async def mock_generator():
yield AssistantMessage(
content=[TextBlock(text="4")]
)
yield AssistantMessage(content=[TextBlock(text="4")])
mock_process.return_value = mock_generator()
@ -35,12 +36,14 @@ class TestQueryFunction:
def test_query_with_options(self):
"""Test query with various options."""
async def _test():
with patch('claude_code_sdk._internal.client.InternalClient.process_query') as mock_process:
with patch(
"claude_code_sdk._internal.client.InternalClient.process_query"
) as mock_process:
async def mock_generator():
yield AssistantMessage(
content=[TextBlock(text="Hello!")]
)
yield AssistantMessage(content=[TextBlock(text="Hello!")])
mock_process.return_value = mock_generator()
@ -48,28 +51,28 @@ class TestQueryFunction:
allowed_tools=["Read", "Write"],
system_prompt="You are helpful",
permission_mode="acceptEdits",
max_turns=5
max_turns=5,
)
messages = []
async for msg in query(
prompt="Hi",
options=options
):
async for msg in query(prompt="Hi", options=options):
messages.append(msg)
# Verify process_query was called with correct prompt and options
mock_process.assert_called_once()
call_args = mock_process.call_args
assert call_args[1]['prompt'] == "Hi"
assert call_args[1]['options'] == options
assert call_args[1]["prompt"] == "Hi"
assert call_args[1]["options"] == options
anyio.run(_test)
def test_query_with_cwd(self):
"""Test query with custom working directory."""
async def _test():
with patch('claude_code_sdk._internal.client.SubprocessCLITransport') as mock_transport_class:
with patch(
"claude_code_sdk._internal.client.SubprocessCLITransport"
) as mock_transport_class:
mock_transport = AsyncMock()
mock_transport_class.return_value = mock_transport
@ -79,8 +82,8 @@ class TestQueryFunction:
"type": "assistant",
"message": {
"role": "assistant",
"content": [{"type": "text", "text": "Done"}]
}
"content": [{"type": "text", "text": "Done"}],
},
}
yield {
"type": "result",
@ -91,7 +94,7 @@ class TestQueryFunction:
"is_error": False,
"num_turns": 1,
"session_id": "test-session",
"total_cost": 0.001
"total_cost": 0.001,
}
mock_transport.receive_messages = mock_receive
@ -100,16 +103,13 @@ class TestQueryFunction:
options = ClaudeCodeOptions(cwd="/custom/path")
messages = []
async for msg in query(
prompt="test",
options=options
):
async for msg in query(prompt="test", options=options):
messages.append(msg)
# Verify transport was created with correct parameters
mock_transport_class.assert_called_once()
call_kwargs = mock_transport_class.call_args.kwargs
assert call_kwargs['prompt'] == "test"
assert call_kwargs['options'].cwd == "/custom/path"
assert call_kwargs["prompt"] == "test"
assert call_kwargs["options"].cwd == "/custom/path"
anyio.run(_test)

View file

@ -23,8 +23,11 @@ class TestIntegration:
def test_simple_query_response(self):
"""Test a simple query with text response."""
async def _test():
with patch("claude_code_sdk._internal.client.SubprocessCLITransport") as mock_transport_class:
with patch(
"claude_code_sdk._internal.client.SubprocessCLITransport"
) as mock_transport_class:
mock_transport = AsyncMock()
mock_transport_class.return_value = mock_transport
@ -71,12 +74,15 @@ class TestIntegration:
assert messages[1].cost_usd == 0.001
assert messages[1].session_id == "test-session"
anyio.run(_test)
def test_query_with_tool_use(self):
"""Test query that uses tools."""
async def _test():
with patch("claude_code_sdk._internal.client.SubprocessCLITransport") as mock_transport_class:
with patch(
"claude_code_sdk._internal.client.SubprocessCLITransport"
) as mock_transport_class:
mock_transport = AsyncMock()
mock_transport_class.return_value = mock_transport
@ -135,25 +141,31 @@ class TestIntegration:
assert messages[0].content[1].name == "Read"
assert messages[0].content[1].input["file_path"] == "/test.txt"
anyio.run(_test)
def test_cli_not_found(self):
"""Test handling when CLI is not found."""
async def _test():
with patch("shutil.which", return_value=None), patch(
"pathlib.Path.exists", return_value=False
), pytest.raises(CLINotFoundError) as exc_info:
with (
patch("shutil.which", return_value=None),
patch("pathlib.Path.exists", return_value=False),
pytest.raises(CLINotFoundError) as exc_info,
):
async for _ in query(prompt="test"):
pass
assert "Claude Code requires Node.js" in str(exc_info.value)
anyio.run(_test)
def test_continuation_option(self):
"""Test query with continue_conversation option."""
async def _test():
with patch("claude_code_sdk._internal.client.SubprocessCLITransport") as mock_transport_class:
with patch(
"claude_code_sdk._internal.client.SubprocessCLITransport"
) as mock_transport_class:
mock_transport = AsyncMock()
mock_transport_class.return_value = mock_transport
@ -179,13 +191,14 @@ class TestIntegration:
# Run query with continuation
messages = []
async for msg in query(
prompt="Continue", options=ClaudeCodeOptions(continue_conversation=True)
prompt="Continue",
options=ClaudeCodeOptions(continue_conversation=True),
):
messages.append(msg)
# Verify transport was created with continuation option
mock_transport_class.assert_called_once()
call_kwargs = mock_transport_class.call_args.kwargs
assert call_kwargs['options'].continue_conversation is True
assert call_kwargs["options"].continue_conversation is True
anyio.run(_test)

View file

@ -16,12 +16,12 @@ class TestSubprocessCLITransport:
"""Test CLI not found error."""
from claude_code_sdk._errors import CLINotFoundError
with patch("shutil.which", return_value=None), patch(
"pathlib.Path.exists", return_value=False
), pytest.raises(CLINotFoundError) as exc_info:
SubprocessCLITransport(
prompt="test", options=ClaudeCodeOptions()
)
with (
patch("shutil.which", return_value=None),
patch("pathlib.Path.exists", return_value=False),
pytest.raises(CLINotFoundError) as exc_info,
):
SubprocessCLITransport(prompt="test", options=ClaudeCodeOptions())
assert "Claude Code requires Node.js" in str(exc_info.value)
@ -43,7 +43,9 @@ class TestSubprocessCLITransport:
from pathlib import Path
transport = SubprocessCLITransport(
prompt="Hello", options=ClaudeCodeOptions(), cli_path=Path("/usr/bin/claude")
prompt="Hello",
options=ClaudeCodeOptions(),
cli_path=Path("/usr/bin/claude"),
)
assert transport._cli_path == "/usr/bin/claude"
@ -92,6 +94,7 @@ class TestSubprocessCLITransport:
def test_connect_disconnect(self):
"""Test connect and disconnect lifecycle."""
async def _test():
with patch("anyio.open_process") as mock_exec:
mock_process = MagicMock()
@ -103,7 +106,9 @@ class TestSubprocessCLITransport:
mock_exec.return_value = mock_process
transport = SubprocessCLITransport(
prompt="test", options=ClaudeCodeOptions(), cli_path="/usr/bin/claude"
prompt="test",
options=ClaudeCodeOptions(),
cli_path="/usr/bin/claude",
)
await transport.connect()

View file

@ -26,9 +26,7 @@ class TestMessageTypes:
def test_tool_use_block(self):
"""Test creating a ToolUseBlock."""
block = ToolUseBlock(
id="tool-123",
name="Read",
input={"file_path": "/test.txt"}
id="tool-123", name="Read", input={"file_path": "/test.txt"}
)
assert block.id == "tool-123"
assert block.name == "Read"
@ -37,9 +35,7 @@ class TestMessageTypes:
def test_tool_result_block(self):
"""Test creating a ToolResultBlock."""
block = ToolResultBlock(
tool_use_id="tool-123",
content="File contents here",
is_error=False
tool_use_id="tool-123", content="File contents here", is_error=False
)
assert block.tool_use_id == "tool-123"
assert block.content == "File contents here"
@ -55,7 +51,7 @@ class TestMessageTypes:
is_error=False,
num_turns=1,
session_id="session-123",
total_cost_usd=0.01
total_cost_usd=0.01,
)
assert msg.subtype == "success"
assert msg.cost_usd == 0.01
@ -78,8 +74,7 @@ class TestOptions:
def test_claude_code_options_with_tools(self):
"""Test Options with built-in tools."""
options = ClaudeCodeOptions(
allowed_tools=["Read", "Write", "Edit"],
disallowed_tools=["Bash"]
allowed_tools=["Read", "Write", "Edit"], disallowed_tools=["Bash"]
)
assert options.allowed_tools == ["Read", "Write", "Edit"]
assert options.disallowed_tools == ["Bash"]
@ -93,25 +88,21 @@ class TestOptions:
"""Test Options with system prompt."""
options = ClaudeCodeOptions(
system_prompt="You are a helpful assistant.",
append_system_prompt="Be concise."
append_system_prompt="Be concise.",
)
assert options.system_prompt == "You are a helpful assistant."
assert options.append_system_prompt == "Be concise."
def test_claude_code_options_with_session_continuation(self):
"""Test Options with session continuation."""
options = ClaudeCodeOptions(
continue_conversation=True,
resume="session-123"
)
options = ClaudeCodeOptions(continue_conversation=True, resume="session-123")
assert options.continue_conversation is True
assert options.resume == "session-123"
def test_claude_code_options_with_model_specification(self):
"""Test Options with model specification."""
options = ClaudeCodeOptions(
model="claude-3-5-sonnet-20241022",
permission_prompt_tool_name="CustomTool"
model="claude-3-5-sonnet-20241022", permission_prompt_tool_name="CustomTool"
)
assert options.model == "claude-3-5-sonnet-20241022"
assert options.permission_prompt_tool_name == "CustomTool"