Improve consistency of convenience imports.

APIs available as convenience imports should have their types also available as convenienc imports. Fix #1560.
python-websockets · Jan 4, 2025 · 2abf77f · 2abf77f
1 parent 197b3ec
commit 2abf77f
Show file tree

Hide file tree

Showing 6 changed files with 77 additions and 29 deletions.
diff --git a/docs/howto/upgrade.rst b/docs/howto/upgrade.rst
@@ -93,8 +93,8 @@ Client APIs
 | ``websockets.client.unix_connect()``                         |br| | :func:`websockets.asyncio.client.unix_connect`      |
 | :func:`websockets.legacy.client.unix_connect`                     |                                                     |
 +-------------------------------------------------------------------+-----------------------------------------------------+
-| ``websockets.WebSocketClientProtocol``                       |br| | :class:`websockets.asyncio.client.ClientConnection` |
-| ``websockets.client.WebSocketClientProtocol``                |br| |                                                     |
+| ``websockets.WebSocketClientProtocol``                       |br| | ``websockets.ClientConnection`` *(since 14.2)* |br| |
+| ``websockets.client.WebSocketClientProtocol``                |br| | :class:`websockets.asyncio.client.ClientConnection` |
 | :class:`websockets.legacy.client.WebSocketClientProtocol`         |                                                     |
 +-------------------------------------------------------------------+-----------------------------------------------------+
 
@@ -112,12 +112,12 @@ Server APIs
 | ``websockets.server.unix_serve()``                           |br| | :func:`websockets.asyncio.server.unix_serve`        |
 | :func:`websockets.legacy.server.unix_serve`                       |                                                     |
 +-------------------------------------------------------------------+-----------------------------------------------------+
-| ``websockets.WebSocketServer``                               |br| | :class:`websockets.asyncio.server.Server`           |
-| ``websockets.server.WebSocketServer``                        |br| |                                                     |
+| ``websockets.WebSocketServer``                               |br| | ``websockets.Server`` *(since 14.2)*           |br| |
+| ``websockets.server.WebSocketServer``                        |br| | :class:`websockets.asyncio.server.Server`           |
 | :class:`websockets.legacy.server.WebSocketServer`                 |                                                     |
 +-------------------------------------------------------------------+-----------------------------------------------------+
-| ``websockets.WebSocketServerProtocol``                       |br| | :class:`websockets.asyncio.server.ServerConnection` |
-| ``websockets.server.WebSocketServerProtocol``                |br| |                                                     |
+| ``websockets.WebSocketServerProtocol``                       |br| | ``websockets.ServerConnection`` *(since 14.2)* |br| |
+| ``websockets.server.WebSocketServerProtocol``                |br| | :class:`websockets.asyncio.server.ServerConnection` |
 | :class:`websockets.legacy.server.WebSocketServerProtocol`         |                                                     |
 +-------------------------------------------------------------------+-----------------------------------------------------+
 | ``websockets.broadcast()`` *(before 14.0)*                   |br| | ``websockets.broadcast()`` *(since 14.0)*      |br| |

diff --git a/src/websockets/__init__.py b/src/websockets/__init__.py
@@ -10,11 +10,14 @@
     # .asyncio.client
     "connect",
     "unix_connect",
+    "ClientConnection",
     # .asyncio.server
     "basic_auth",
     "broadcast",
     "serve",
     "unix_serve",
+    "ServerConnection",
+    "Server",
     # .client
     "ClientProtocol",
     # .datastructures
@@ -44,6 +47,18 @@
     "ProtocolError",
     "SecurityError",
     "WebSocketException",
+    # .frames
+    "Close",
+    "CloseCode",
+    "Frame",
+    "Opcode",
+    # .http11
+    "Request",
+    "Response",
+    # .protocol
+    "Protocol",
+    "Side",
+    "State",
     # .server
     "ServerProtocol",
     # .typing
@@ -58,8 +73,15 @@
 
 # When type checking, import non-deprecated aliases eagerly. Else, import on demand.
 if typing.TYPE_CHECKING:
-    from .asyncio.client import connect, unix_connect
-    from .asyncio.server import basic_auth, broadcast, serve, unix_serve
+    from .asyncio.client import ClientConnection, connect, unix_connect
+    from .asyncio.server import (
+        Server,
+        ServerConnection,
+        basic_auth,
+        broadcast,
+        serve,
+        unix_serve,
+    )
     from .client import ClientProtocol
     from .datastructures import Headers, HeadersLike, MultipleValuesError
     from .exceptions import (
@@ -86,6 +108,9 @@
         SecurityError,
         WebSocketException,
     )
+    from .frames import Close, CloseCode, Frame, Opcode
+    from .http11 import Request, Response
+    from .protocol import Protocol, Side, State
     from .server import ServerProtocol
     from .typing import (
         Data,
@@ -103,11 +128,14 @@
             # .asyncio.client
             "connect": ".asyncio.client",
             "unix_connect": ".asyncio.client",
+            "ClientConnection": ".asyncio.client",
             # .asyncio.server
             "basic_auth": ".asyncio.server",
             "broadcast": ".asyncio.server",
             "serve": ".asyncio.server",
             "unix_serve": ".asyncio.server",
+            "ServerConnection": ".asyncio.server",
+            "Server": ".asyncio.server",
             # .client
             "ClientProtocol": ".client",
             # .datastructures
@@ -137,6 +165,18 @@
             "ProtocolError": ".exceptions",
             "SecurityError": ".exceptions",
             "WebSocketException": ".exceptions",
+            # .frames
+            "Close": ".frames",
+            "CloseCode": ".frames",
+            "Frame": ".frames",
+            "Opcode": ".frames",
+            # .http11
+            "Request": ".http11",
+            "Response": ".http11",
+            # .protocol
+            "Protocol": ".protocol",
+            "Side": ".protocol",
+            "State": ".protocol",
             # .server
             "ServerProtocol": ".server",
             # .typing

diff --git a/src/websockets/datastructures.py b/src/websockets/datastructures.py
@@ -4,7 +4,11 @@
 from typing import Any, Protocol, Union
 
 
-__all__ = ["Headers", "HeadersLike", "MultipleValuesError"]
+__all__ = [
+    "Headers",
+    "HeadersLike",
+    "MultipleValuesError",
+]
 
 
 class MultipleValuesError(LookupError):

diff --git a/src/websockets/frames.py b/src/websockets/frames.py
@@ -28,6 +28,7 @@
     "OP_PONG",
     "DATA_OPCODES",
     "CTRL_OPCODES",
+    "CloseCode",
     "Frame",
     "Close",
 ]

diff --git a/src/websockets/http11.py b/src/websockets/http11.py
@@ -13,7 +13,12 @@
 from .version import version as websockets_version
 
 
-__all__ = ["SERVER", "USER_AGENT", "Request", "Response"]
+__all__ = [
+    "SERVER",
+    "USER_AGENT",
+    "Request",
+    "Response",
+]
 
 
 PYTHON_VERSION = "{}.{}".format(*sys.version_info)

diff --git a/tests/test_exports.py b/tests/test_exports.py
@@ -11,32 +11,30 @@
 import websockets.uri
 
 
-combined_exports = (
-    []
-    + websockets.asyncio.client.__all__
-    + websockets.asyncio.server.__all__
-    + websockets.client.__all__
-    + websockets.datastructures.__all__
-    + websockets.exceptions.__all__
-    + websockets.server.__all__
-    + websockets.typing.__all__
-)
-
-# These API are intentionally not re-exported by the top-level module.
-missing_reexports = [
-    # websockets.asyncio.client
-    "ClientConnection",
-    # websockets.asyncio.server
-    "ServerConnection",
-    "Server",
+combined_exports = [
+    name
+    for name in (
+        []
+        + websockets.asyncio.client.__all__
+        + websockets.asyncio.server.__all__
+        + websockets.client.__all__
+        + websockets.datastructures.__all__
+        + websockets.exceptions.__all__
+        + websockets.frames.__all__
+        + websockets.http11.__all__
+        + websockets.protocol.__all__
+        + websockets.server.__all__
+        + websockets.typing.__all__
+    )
+    if not name.isupper()  # filter out constants
 ]
 
 
 class ExportsTests(unittest.TestCase):
     def test_top_level_module_reexports_submodule_exports(self):
         self.assertEqual(
             set(combined_exports),
-            set(websockets.__all__ + missing_reexports),
+            set(websockets.__all__),
         )
 
     def test_submodule_exports_are_globally_unique(self):