Take advantage of typing.Self.

This is a small correctness improvement for Connection classes, which
are subclassed as ClientConnection and ServerConnection.

Author: aaugustin

src/websockets/asyncio/connection.py

@@ -10,7 +10,7 @@
 import uuid
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Iterable, Mapping
 from types import TracebackType
-from typing import Any, Literal, cast, overload
+from typing import Any, Literal, Self, cast, overload
 
 from ..exceptions import (
     ConcurrencyError,
@@ -205,7 +205,7 @@ def close_reason(self) -> str | None:
 
     # Public methods
 
-    async def __aenter__(self) -> Connection:
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(

src/websockets/asyncio/server.py

@@ -8,7 +8,7 @@
 import socket
 from collections.abc import Awaitable, Generator, Iterable, Sequence
 from types import TracebackType
-from typing import Any, Callable, Mapping, cast
+from typing import Any, Callable, Mapping, Self, cast
 
 from ..exceptions import InvalidHeader
 from ..extensions.base import ServerExtensionFactory
@@ -543,7 +543,7 @@ def sockets(self) -> tuple[socket.socket, ...]:
         """
         return self.server.sockets
 
-    async def __aenter__(self) -> Server:  # pragma: no cover
+    async def __aenter__(self) -> Self:  # pragma: no cover
         return self
 
     async def __aexit__(

src/websockets/frames.py

@@ -7,7 +7,7 @@
 import secrets
 import struct
 from collections.abc import Generator, Sequence
-from typing import Callable
+from typing import Callable, Self
 
 from .exceptions import PayloadTooBig, ProtocolError
 from .typing import BytesLike
@@ -431,7 +431,7 @@ def __str__(self) -> str:
         return result
 
     @classmethod
-    def parse(cls, data: BytesLike) -> Close:
+    def parse(cls, data: BytesLike) -> Self:
         """
         Parse the payload of a close frame.
 

src/websockets/legacy/server.py

@@ -10,7 +10,7 @@
 import warnings
 from collections.abc import Awaitable, Generator, Iterable, Sequence
 from types import TracebackType
-from typing import Any, Callable, cast
+from typing import Any, Callable, Self, cast
 
 from ..datastructures import Headers, HeadersLike, MultipleValuesError
 from ..exceptions import (
@@ -861,7 +861,7 @@ def sockets(self) -> Iterable[socket.socket]:
         """
         return self.server.sockets
 
-    async def __aenter__(self) -> WebSocketServer:  # pragma: no cover
+    async def __aenter__(self) -> Self:  # pragma: no cover
         return self
 
     async def __aexit__(

src/websockets/sync/connection.py

@@ -12,7 +12,7 @@
 import uuid
 from collections.abc import Iterable, Iterator, Mapping
 from types import TracebackType
-from typing import Any, Literal, overload
+from typing import Any, Literal, Self, overload
 
 from ..exceptions import (
     ConcurrencyError,
@@ -216,7 +216,7 @@ def close_reason(self) -> str | None:
 
     # Public methods
 
-    def __enter__(self) -> Connection:
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(

src/websockets/sync/server.py

@@ -14,7 +14,7 @@
 import warnings
 from collections.abc import Iterable, Sequence
 from types import TracebackType
-from typing import Any, Callable, Mapping, cast
+from typing import Any, Callable, Mapping, Self, cast
 
 from ..exceptions import InvalidHeader
 from ..extensions.base import ServerExtensionFactory
@@ -439,7 +439,7 @@ def fileno(self) -> int:
         """
         return self.socket.fileno()
 
-    def __enter__(self) -> Server:
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(

src/websockets/trio/connection.py

@@ -8,7 +8,7 @@
 import uuid
 from collections.abc import AsyncIterable, AsyncIterator, Iterable, Mapping
 from types import TracebackType
-from typing import Any, Literal, overload
+from typing import Any, Literal, Self, overload
 
 import trio
 import trio.abc
@@ -221,7 +221,7 @@ def close_reason(self) -> str | None:
 
     # Public methods
 
-    async def __aenter__(self) -> Connection:
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(

src/websockets/trio/server.py

@@ -7,7 +7,7 @@
 import ssl as ssl_module
 from collections.abc import Awaitable, Mapping, Sequence
 from types import TracebackType
-from typing import Any, Callable
+from typing import Any, Callable, Self
 
 import trio
 import trio.abc
@@ -338,7 +338,7 @@ async def aclose(
 
         self.logger.info("server closed")
 
-    async def __aenter__(self) -> Server:
+    async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(