uniformize context managers

src/websockets/asyncio/client.py

@@ -574,7 +574,7 @@ async def __aenter__(self) -> ClientConnection:
         if hasattr(self, "connection"):
             raise RuntimeError("connect() isn't reentrant")
         self.connection = await self
-        return self.connection
+        return await self.connection.__aenter__()
 
     async def __aexit__(
         self,
@@ -583,7 +583,7 @@ async def __aexit__(
         traceback: TracebackType | None,
     ) -> None:
         try:
-            await self.connection.close()
+            return await self.connection.__aexit__(exc_type, exc_value, traceback)
         finally:
             del self.connection
 

src/websockets/asyncio/server.py

@@ -744,12 +744,10 @@ async def protocol_handler(connection: ServerConnection) -> None:
             server.all_connections.add(connection)
             connection.start_keepalive()
             try:
-                await handler(connection)
+                async with connection:
+                    await handler(connection)
             except Exception:
                 connection.logger.error("connection handler failed", exc_info=True)
-                await connection.close(CloseCode.INTERNAL_ERROR)
-            else:
-                await connection.close()
             finally:
                 server.all_connections.discard(connection)
 

src/websockets/sync/client.py

@@ -521,7 +521,7 @@ def __enter__(self) -> ClientConnection:
             raise RuntimeError("connect() isn't reentrant")
         self.connection = self.connect()
         self.connection.pending_legacy_warning = False
-        return self.connection
+        return self.connection.__enter__()
 
     def __exit__(
         self,
@@ -530,7 +530,7 @@ def __exit__(
         exc_traceback: TracebackType | None,
     ) -> None:
         try:
-            self.connection.close()
+            return self.connection.__exit__(exc_type, exc_value, exc_traceback)
         finally:
             del self.connection
 

src/websockets/sync/server.py

@@ -770,12 +770,10 @@ def protocol_select_subprotocol(
                 server.all_connections.add(connection)
             connection.start_keepalive()
             try:
-                handler(connection)
+                with connection:
+                    handler(connection)
             except Exception:
                 connection.logger.error("connection handler failed", exc_info=True)
-                connection.close(CloseCode.INTERNAL_ERROR)
-            else:
-                connection.close()
             finally:
                 with server.lock:
                     server.all_connections.discard(connection)

src/websockets/trio/client.py

@@ -549,7 +549,7 @@ async def __aenter__(self) -> ClientConnection:
         await self.__aenter_nursery__()
         try:
             self.connection = await self.connect(self.nursery)
-            return self.connection
+            return await self.connection.__aenter__()
         except BaseException as exc:
             await self.__aexit_nursery__(type(exc), exc, exc.__traceback__)
             raise AssertionError("expected __aexit_nursery__ to re-raise the exception")
@@ -562,7 +562,7 @@ async def __aexit__(
     ) -> None:
         try:
             try:
-                await self.connection.aclose()
+                return await self.connection.__aexit__(exc_type, exc_value, traceback)
             finally:
                 del self.connection
         finally:

src/websockets/trio/server.py

@@ -639,12 +639,10 @@ def protocol_select_subprotocol(
                 server.all_connections.add(connection)
                 connection.start_keepalive()
                 try:
-                    await handler(connection)
+                    async with connection:
+                        await handler(connection)
                 except Exception:
                     connection.logger.error("connection handler failed", exc_info=True)
-                    await connection.aclose(CloseCode.INTERNAL_ERROR)
-                else:
-                    await connection.aclose()
                 finally:
                     server.all_connections.discard(connection)
 

tests/asyncio/test_client.py

@@ -60,6 +60,21 @@ async def test_context_manager(self):
                 self.assertEqual(client.protocol.state.name, "OPEN")
             self.assertEqual(client.protocol.state.name, "CLOSED")
 
+    async def test_context_manager_normal_exit(self):
+        """Client closes the connection with code 1000 when exiting normally."""
+        async with serve(*args) as server:
+            async with connect(get_uri(server)) as client:
+                pass
+            self.assertEqual(client.close_code, 1000)
+
+    async def test_context_manager_exception(self):
+        """Client closes the connection with code 1011 when exiting with an error."""
+        async with serve(*args) as server:
+            with self.assertRaises(RuntimeError):
+                async with connect(get_uri(server)) as client:
+                    raise RuntimeError("BOOM")
+            self.assertEqual(client.close_code, 1011)
+
     async def test_direct_connection(self):
         """Client connects to server directly."""
         async with serve(*args) as server:

tests/sync/test_client.py

@@ -67,6 +67,21 @@ def test_context_manager(self):
                 self.assertEqual(client.protocol.state.name, "OPEN")
             self.assertEqual(client.protocol.state.name, "CLOSED")
 
+    def test_context_manager_normal_exit(self):
+        """Client closes the connection with code 1000 when exiting normally."""
+        with run_server() as server:
+            with connect(get_uri(server)) as client:
+                pass
+            self.assertEqual(client.close_code, 1000)
+
+    def test_context_manager_exception(self):
+        """Client closes the connection with code 1011 when exiting with an error."""
+        with run_server() as server:
+            with self.assertRaises(RuntimeError):
+                with connect(get_uri(server)) as client:
+                    raise RuntimeError("BOOM")
+            self.assertEqual(client.close_code, 1011)
+
     def test_direct_connection(self):
         """Client connects to server directly."""
         with run_server() as server:

tests/trio/test_client.py

@@ -62,6 +62,21 @@ async def test_context_manager(self):
                 self.assertEqual(client.protocol.state.name, "OPEN")
             self.assertEqual(client.protocol.state.name, "CLOSED")
 
+    async def test_context_manager_normal_exit(self):
+        """Client closes the connection with code 1000 when exiting normally."""
+        async with run_server() as server:
+            async with connect(get_uri(server)) as client:
+                pass
+            self.assertEqual(client.close_code, 1000)
+
+    async def test_context_manager_exception(self):
+        """Client closes the connection with code 1011 when exiting with an error."""
+        async with run_server() as server:
+            with self.assertRaises(RuntimeError):
+                async with connect(get_uri(server)) as client:
+                    raise RuntimeError("BOOM")
+            self.assertEqual(client.close_code, 1011)
+
     async def test_explicit_host_port(self):
         """Client connects using an explicit host / port."""
         async with run_server() as server: