Format with ruff

dsp-ant · dsp-ant · commit 1657212d6836 · 2024-10-11T16:07:27.000+01:00
diff --git a/mcp_python/__init__.py b/mcp_python/__init__.py
@@ -37,7 +37,6 @@
     ReadResourceResult,
     Resource,
     ResourceUpdatedNotification,
-    Role as SamplingRole,
     SamplingMessage,
     ServerCapabilities,
     ServerNotification,
@@ -49,6 +48,9 @@
     Tool,
     UnsubscribeRequest,
 )
+from .types import (
+    Role as SamplingRole,
+)
 
 __all__ = [
     "CallToolRequest",
diff --git a/mcp_python/client/session.py b/mcp_python/client/session.py
@@ -62,7 +62,8 @@ async def initialize(self) -> InitializeResult:
 
         if result.protocolVersion != SUPPORTED_PROTOCOL_VERSION:
             raise RuntimeError(
-                f"Unsupported protocol version from the server: {result.protocolVersion}"
+                "Unsupported protocol version from the server: "
+                f"{result.protocolVersion}"
             )
 
         await self.send_notification(
diff --git a/mcp_python/client/sse.py b/mcp_python/client/sse.py
@@ -19,11 +19,17 @@ def remove_request_params(url: str) -> str:
 
 
 @asynccontextmanager
-async def sse_client(url: str, headers: dict[str, Any] | None = None, timeout: float = 5, sse_read_timeout: float = 60 * 5):
+async def sse_client(
+    url: str,
+    headers: dict[str, Any] | None = None,
+    timeout: float = 5,
+    sse_read_timeout: float = 60 * 5,
+):
     """
     Client transport for SSE.
 
-    `sse_read_timeout` determines how long (in seconds) the client will wait for a new event before disconnecting. All other HTTP operations are controlled by `timeout`.
+    `sse_read_timeout` determines how long (in seconds) the client will wait for a new
+    event before disconnecting. All other HTTP operations are controlled by `timeout`.
     """
     read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
     read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
@@ -67,7 +73,10 @@ async def sse_reader(
                                             or url_parsed.scheme
                                             != endpoint_parsed.scheme
                                         ):
-                                            error_msg = f"Endpoint origin does not match connection origin: {endpoint_url}"
+                                            error_msg = (
+                                                "Endpoint origin does not match "
+                                                f"connection origin: {endpoint_url}"
+                                            )
                                             logger.error(error_msg)
                                             raise ValueError(error_msg)
 
@@ -104,11 +113,16 @@ async def post_writer(endpoint_url: str):
                                     logger.debug(f"Sending client message: {message}")
                                     response = await client.post(
                                         endpoint_url,
-                                        json=message.model_dump(by_alias=True, mode="json", exclude_none=True),
+                                        json=message.model_dump(
+                                            by_alias=True,
+                                            mode="json",
+                                            exclude_none=True,
+                                        ),
                                     )
                                     response.raise_for_status()
                                     logger.debug(
-                                        f"Client message sent successfully: {response.status_code}"
+                                        "Client message sent successfully: "
+                                        f"{response.status_code}"
                                     )
                         except Exception as exc:
                             logger.error(f"Error in post_writer: {exc}")
diff --git a/mcp_python/client/stdio.py b/mcp_python/client/stdio.py
@@ -28,7 +28,8 @@ class StdioServerParameters(BaseModel):
 @asynccontextmanager
 async def stdio_client(server: StdioServerParameters):
     """
-    Client transport for stdio: this will connect to a server by spawning a process and communicating with it over stdin/stdout.
+    Client transport for stdio: this will connect to a server by spawning a
+    process and communicating with it over stdin/stdout.
     """
     read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception]
     read_stream_writer: MemoryObjectSendStream[JSONRPCMessage | Exception]
diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py
@@ -45,7 +45,6 @@
     "request_ctx"
 )
 
-
 class Server:
     def __init__(self, name: str):
         self.name = name
@@ -87,7 +86,7 @@ def request_context(self) -> RequestContext:
 
     def list_prompts(self):
         def decorator(func: Callable[[], Awaitable[list[Prompt]]]):
-            logger.debug(f"Registering handler for PromptListRequest")
+            logger.debug("Registering handler for PromptListRequest")
 
             async def handler(_: Any):
                 prompts = await func()
@@ -103,17 +102,19 @@ def get_prompt(self):
             GetPromptRequest,
             GetPromptResult,
             ImageContent,
-            Role as Role,
             SamplingMessage,
             TextContent,
         )
+        from mcp_python.types import (
+            Role as Role,
+        )
 
         def decorator(
             func: Callable[
                 [str, dict[str, str] | None], Awaitable[types.PromptResponse]
             ],
         ):
-            logger.debug(f"Registering handler for GetPromptRequest")
+            logger.debug("Registering handler for GetPromptRequest")
 
             async def handler(req: GetPromptRequest):
                 prompt_get = await func(req.params.name, req.params.arguments)
@@ -149,7 +150,7 @@ async def handler(req: GetPromptRequest):
 
     def list_resources(self):
         def decorator(func: Callable[[], Awaitable[list[Resource]]]):
-            logger.debug(f"Registering handler for ListResourcesRequest")
+            logger.debug("Registering handler for ListResourcesRequest")
 
             async def handler(_: Any):
                 resources = await func()
@@ -169,7 +170,7 @@ def read_resource(self):
         )
 
         def decorator(func: Callable[[AnyUrl], Awaitable[str | bytes]]):
-            logger.debug(f"Registering handler for ReadResourceRequest")
+            logger.debug("Registering handler for ReadResourceRequest")
 
             async def handler(req: ReadResourceRequest):
                 result = await func(req.params.uri)
@@ -204,7 +205,7 @@ def set_logging_level(self):
         from mcp_python.types import EmptyResult
 
         def decorator(func: Callable[[LoggingLevel], Awaitable[None]]):
-            logger.debug(f"Registering handler for SetLevelRequest")
+            logger.debug("Registering handler for SetLevelRequest")
 
             async def handler(req: SetLevelRequest):
                 await func(req.params.level)
@@ -219,7 +220,7 @@ def subscribe_resource(self):
         from mcp_python.types import EmptyResult
 
         def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
-            logger.debug(f"Registering handler for SubscribeRequest")
+            logger.debug("Registering handler for SubscribeRequest")
 
             async def handler(req: SubscribeRequest):
                 await func(req.params.uri)
@@ -234,7 +235,7 @@ def unsubscribe_resource(self):
         from mcp_python.types import EmptyResult
 
         def decorator(func: Callable[[AnyUrl], Awaitable[None]]):
-            logger.debug(f"Registering handler for UnsubscribeRequest")
+            logger.debug("Registering handler for UnsubscribeRequest")
 
             async def handler(req: UnsubscribeRequest):
                 await func(req.params.uri)
@@ -249,7 +250,7 @@ def call_tool(self):
         from mcp_python.types import CallToolResult
 
         def decorator(func: Callable[..., Awaitable[Any]]):
-            logger.debug(f"Registering handler for CallToolRequest")
+            logger.debug("Registering handler for CallToolRequest")
 
             async def handler(req: CallToolRequest):
                 result = await func(req.params.name, **(req.params.arguments or {}))
@@ -264,7 +265,7 @@ def progress_notification(self):
         def decorator(
             func: Callable[[str | int, float, float | None], Awaitable[None]],
         ):
-            logger.debug(f"Registering handler for ProgressNotification")
+            logger.debug("Registering handler for ProgressNotification")
 
             async def handler(req: ProgressNotification):
                 await func(
@@ -286,7 +287,7 @@ def decorator(
                 Awaitable[Completion | None],
             ],
         ):
-            logger.debug(f"Registering handler for CompleteRequest")
+            logger.debug("Registering handler for CompleteRequest")
 
             async def handler(req: CompleteRequest):
                 completion = await func(req.params.ref, req.params.argument)
@@ -310,7 +311,9 @@ async def run(
         initialization_options: types.InitializationOptions
     ):
         with warnings.catch_warnings(record=True) as w:
-            async with ServerSession(read_stream, write_stream, initialization_options) as session:
+            async with ServerSession(
+                read_stream, write_stream, initialization_options
+            ) as session:
                 async for message in session.incoming_messages:
                     logger.debug(f"Received message: {message}")
 
@@ -359,14 +362,16 @@ async def run(
 
                                 handler = self.notification_handlers[type(notify)]
                                 logger.debug(
-                                    f"Dispatching notification of type {type(notify).__name__}"
+                                    f"Dispatching notification of type "
+                                    f"{type(notify).__name__}"
                                 )
 
                                 try:
                                     await handler(notify)
                                 except Exception as err:
                                     logger.error(
-                                        f"Uncaught exception in notification handler: {err}"
+                                        f"Uncaught exception in notification handler: "
+                                        f"{err}"
                                     )
 
                     for warning in w:
diff --git a/mcp_python/server/__main__.py b/mcp_python/server/__main__.py
@@ -1,6 +1,7 @@
+import importlib.metadata
 import logging
 import sys
-import importlib.metadata
+
 import anyio
 
 from mcp_python.server.session import ServerSession
@@ -30,7 +31,18 @@ async def receive_loop(session: ServerSession):
 async def main():
     version = importlib.metadata.version("mcp_python")
     async with stdio_server() as (read_stream, write_stream):
-        async with ServerSession(read_stream, write_stream, InitializationOptions(server_name="mcp_python", server_version=version, capabilities=ServerCapabilities())) as session, write_stream:
+        async with (
+            ServerSession(
+                read_stream,
+                write_stream,
+                InitializationOptions(
+                    server_name="mcp_python",
+                    server_version=version,
+                    capabilities=ServerCapabilities(),
+                ),
+            ) as session,
+            write_stream,
+        ):
             await receive_loop(session)
 
 
diff --git a/mcp_python/server/session.py b/mcp_python/server/session.py
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from enum import Enum
 from typing import Any
 
@@ -72,7 +73,7 @@ async def _received_request(
                             capabilities=self._init_options.capabilities,
                             serverInfo=Implementation(
                                 name=self._init_options.server_name,
-                                version=self._init_options.server_version
+                                version=self._init_options.server_version,
                             ),
                         )
                     )
diff --git a/mcp_python/server/sse.py b/mcp_python/server/sse.py
@@ -19,18 +19,23 @@
 
 class SseServerTransport:
     """
-    SSE server transport for MCP. This class provides _two_ ASGI applications, suitable to be used with a framework like Starlette and a server like Hypercorn:
-
-        1. connect_sse() is an ASGI application which receives incoming GET requests, and sets up a new SSE stream to send server messages to the client.
-        2. handle_post_message() is an ASGI application which receives incoming POST requests, which should contain client messages that link to a previously-established SSE session.
+    SSE server transport for MCP. This class provides _two_ ASGI applications,
+    suitable to be used with a framework like Starlette and a server like Hypercorn:
+
+        1. connect_sse() is an ASGI application which receives incoming GET requests,
+           and sets up a new SSE stream to send server messages to the client.
+        2. handle_post_message() is an ASGI application which receives incoming POST
+           requests, which should contain client messages that link to a
+           previously-established SSE session.
     """
 
     _endpoint: str
     _read_stream_writers: dict[UUID, MemoryObjectSendStream[JSONRPCMessage | Exception]]
 
     def __init__(self, endpoint: str) -> None:
         """
-        Creates a new SSE server transport, which will direct the client to POST messages to the relative or absolute URL given.
+        Creates a new SSE server transport, which will direct the client to POST
+        messages to the relative or absolute URL given.
         """
 
         super().__init__()
@@ -74,7 +79,9 @@ async def sse_writer():
                     await sse_stream_writer.send(
                         {
                             "event": "message",
-                            "data": message.model_dump_json(by_alias=True, exclude_none=True),
+                            "data": message.model_dump_json(
+                                by_alias=True, exclude_none=True
+                            ),
                         }
                     )
 
diff --git a/mcp_python/server/stdio.py b/mcp_python/server/stdio.py
@@ -7,14 +7,18 @@
 
 from mcp_python.types import JSONRPCMessage
 
+
 @asynccontextmanager
 async def stdio_server(
-    stdin: anyio.AsyncFile[str] | None = None, stdout: anyio.AsyncFile[str] | None = None
+    stdin: anyio.AsyncFile[str] | None = None,
+    stdout: anyio.AsyncFile[str] | None = None,
 ):
     """
-    Server transport for stdio: this communicates with an MCP client by reading from the current process' stdin and writing to stdout.
+    Server transport for stdio: this communicates with an MCP client by reading
+    from the current process' stdin and writing to stdout.
     """
-    # Purposely not using context managers for these, as we don't want to close standard process handles.
+    # Purposely not using context managers for these, as we don't want to close
+    # standard process handles.
     if not stdin:
         stdin = anyio.wrap_file(sys.stdin)
     if not stdout:
diff --git a/mcp_python/server/websocket.py b/mcp_python/server/websocket.py
@@ -14,7 +14,8 @@
 @asynccontextmanager
 async def websocket_server(scope: Scope, receive: Receive, send: Send):
     """
-    WebSocket server transport for MCP. This is an ASGI application, suitable to be used with a framework like Starlette and a server like Hypercorn.
+    WebSocket server transport for MCP. This is an ASGI application, suitable to be
+    used with a framework like Starlette and a server like Hypercorn.
     """
 
     websocket = WebSocket(scope, receive, send)
@@ -47,7 +48,9 @@ async def ws_writer():
         try:
             async with write_stream_reader:
                 async for message in write_stream_reader:
-                    obj = message.model_dump(by_alias=True, mode="json", exclude_none=True)
+                    obj = message.model_dump(
+                        by_alias=True, mode="json", exclude_none=True
+                    )
                     await websocket.send_json(obj)
         except anyio.ClosedResourceError:
             await websocket.close()
diff --git a/mcp_python/shared/session.py b/mcp_python/shared/session.py
diff --git a/mcp_python/types.py b/mcp_python/types.py
diff --git a/tests/client/test_session.py b/tests/client/test_session.py
diff --git a/tests/server/test_session.py b/tests/server/test_session.py

Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,8 @@ async def initialize(self) -> InitializeResult:`
`62`	`62`
`63`	`63`	`if result.protocolVersion != SUPPORTED_PROTOCOL_VERSION:`
`64`	`64`	`raise RuntimeError(`
`65`		`- f"Unsupported protocol version from the server: {result.protocolVersion}"`
	`65`	`+ "Unsupported protocol version from the server: "`
	`66`	`+ f"{result.protocolVersion}"`
`66`	`67`	`)`
`67`	`68`
`68`	`69`	`await self.send_notification(`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+from dataclasses import dataclass`
`1`	`2`	`from enum import Enum`
`2`	`3`	`from typing import Any`
`3`	`4`
`@@ -72,7 +73,7 @@ async def _received_request(`
`72`	`73`	`capabilities=self._init_options.capabilities,`
`73`	`74`	`serverInfo=Implementation(`
`74`	`75`	`name=self._init_options.server_name,`
`75`		`- version=self._init_options.server_version`
	`76`	`+ version=self._init_options.server_version,`
`76`	`77`	`),`
`77`	`78`	`)`
`78`	`79`	`)`