Import eagerly when running under a type checker.

Fix #1292 (and many others). Also fix some inconsistencies in the lists. The new rule is: - when re-exporting a module, re-export it entirely; - don't re-export deprecated modules.
python-websockets · Oct 1, 2023 · 439dafa · 439dafa
1 parent ca5926e
commit 439dafa
Show file tree

Hide file tree

Showing 6 changed files with 175 additions and 127 deletions.
diff --git a/docs/faq/misc.rst b/docs/faq/misc.rst
@@ -12,29 +12,6 @@ instead of the websockets library.
 
 .. _real-import-paths:
 
-Why does my IDE fail to show documentation for websockets APIs?
-...............................................................
-
-You are probably using the convenience imports e.g.::
-
-    import websockets
-
-    websockets.connect(...)
-    websockets.serve(...)
-
-This is incompatible with static code analysis. It may break auto-completion and
-contextual documentation in IDEs, type checking with mypy_, etc.
-
-.. _mypy: https://github.com/python/mypy
-
-Instead, use the real import paths e.g.::
-
-    import websockets.client
-    import websockets.server
-
-    websockets.client.connect(...)
-    websockets.server.serve(...)
-
 Why is the default implementation located in ``websockets.legacy``?
 ...................................................................
 

diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst
@@ -41,6 +41,11 @@ Backwards-incompatible changes
 Improvements
 ............
 
+* Made convenience imports from ``websockets`` compatible with static code
+  analysis tools such as auto-completion in an IDE or type checking with mypy_.
+
+  .. _mypy: https://github.com/python/mypy
+
 * Added :class:`~frames.CloseCode`.
 
 11.0.3

diff --git a/docs/reference/index.rst b/docs/reference/index.rst
@@ -88,16 +88,3 @@ Convenience imports
 
 For convenience, many public APIs can be imported directly from the
 ``websockets`` package.
-
-
-.. admonition:: Convenience imports are incompatible with some development tools.
-    :class: caution
-
-    Specifically, static code analysis tools don't understand them. This breaks
-    auto-completion and contextual documentation in IDEs, type checking with
-    mypy_, etc.
-
-    .. _mypy: https://github.com/python/mypy
-
-    If you're using such tools, stick to the full import paths, as explained in
-    this FAQ: :ref:`real-import-paths`
diff --git a/src/websockets/__init__.py b/src/websockets/__init__.py
@@ -1,23 +1,24 @@
 from __future__ import annotations
 
+import typing
+
 from .imports import lazy_import
 from .version import version as __version__  # noqa: F401
 
 
 __all__ = [
-    "AbortHandshake",
-    "basic_auth_protocol_factory",
-    "BasicAuthWebSocketServerProtocol",
-    "broadcast",
+    # .client
     "ClientProtocol",
-    "connect",
+    # .datastructures
+    "Headers",
+    "HeadersLike",
+    "MultipleValuesError",
+    # .exceptions
+    "AbortHandshake",
     "ConnectionClosed",
     "ConnectionClosedError",
     "ConnectionClosedOK",
-    "Data",
     "DuplicateParameter",
-    "ExtensionName",
-    "ExtensionParameter",
     "InvalidHandshake",
     "InvalidHeader",
     "InvalidHeaderFormat",
@@ -31,84 +32,156 @@
     "InvalidStatusCode",
     "InvalidUpgrade",
     "InvalidURI",
-    "LoggerLike",
     "NegotiationError",
-    "Origin",
-    "parse_uri",
     "PayloadTooBig",
     "ProtocolError",
     "RedirectHandshake",
     "SecurityError",
-    "serve",
-    "ServerProtocol",
-    "Subprotocol",
-    "unix_connect",
-    "unix_serve",
-    "WebSocketClientProtocol",
-    "WebSocketCommonProtocol",
     "WebSocketException",
     "WebSocketProtocolError",
+    # .legacy.auth
+    "BasicAuthWebSocketServerProtocol",
+    "basic_auth_protocol_factory",
+    # .legacy.client
+    "WebSocketClientProtocol",
+    "connect",
+    "unix_connect",
+    # .legacy.protocol
+    "WebSocketCommonProtocol",
+    "broadcast",
+    # .legacy.server
     "WebSocketServer",
     "WebSocketServerProtocol",
-    "WebSocketURI",
+    "serve",
+    "unix_serve",
+    # .server
+    "ServerProtocol",
+    # .typing
+    "Data",
+    "ExtensionName",
+    "ExtensionParameter",
+    "LoggerLike",
+    "Origin",
+    "Subprotocol",
 ]
 
-lazy_import(
-    globals(),
-    aliases={
-        "auth": ".legacy",
-        "basic_auth_protocol_factory": ".legacy.auth",
-        "BasicAuthWebSocketServerProtocol": ".legacy.auth",
-        "broadcast": ".legacy.protocol",
-        "ClientProtocol": ".client",
-        "connect": ".legacy.client",
-        "unix_connect": ".legacy.client",
-        "WebSocketClientProtocol": ".legacy.client",
-        "Headers": ".datastructures",
-        "MultipleValuesError": ".datastructures",
-        "WebSocketException": ".exceptions",
-        "ConnectionClosed": ".exceptions",
-        "ConnectionClosedError": ".exceptions",
-        "ConnectionClosedOK": ".exceptions",
-        "InvalidHandshake": ".exceptions",
-        "SecurityError": ".exceptions",
-        "InvalidMessage": ".exceptions",
-        "InvalidHeader": ".exceptions",
-        "InvalidHeaderFormat": ".exceptions",
-        "InvalidHeaderValue": ".exceptions",
-        "InvalidOrigin": ".exceptions",
-        "InvalidUpgrade": ".exceptions",
-        "InvalidStatus": ".exceptions",
-        "InvalidStatusCode": ".exceptions",
-        "NegotiationError": ".exceptions",
-        "DuplicateParameter": ".exceptions",
-        "InvalidParameterName": ".exceptions",
-        "InvalidParameterValue": ".exceptions",
-        "AbortHandshake": ".exceptions",
-        "RedirectHandshake": ".exceptions",
-        "InvalidState": ".exceptions",
-        "InvalidURI": ".exceptions",
-        "PayloadTooBig": ".exceptions",
-        "ProtocolError": ".exceptions",
-        "WebSocketProtocolError": ".exceptions",
-        "protocol": ".legacy",
-        "WebSocketCommonProtocol": ".legacy.protocol",
-        "ServerProtocol": ".server",
-        "serve": ".legacy.server",
-        "unix_serve": ".legacy.server",
-        "WebSocketServerProtocol": ".legacy.server",
-        "WebSocketServer": ".legacy.server",
-        "Data": ".typing",
-        "LoggerLike": ".typing",
-        "Origin": ".typing",
-        "ExtensionHeader": ".typing",
-        "ExtensionParameter": ".typing",
-        "Subprotocol": ".typing",
-    },
-    deprecated_aliases={
-        "framing": ".legacy",
-        "handshake": ".legacy",
-        "parse_uri": ".uri",
-        "WebSocketURI": ".uri",
-    },
-)
+# When type checking, import non-deprecated aliases eagerly. Else, import on demand.
+if typing.TYPE_CHECKING:
+    from .client import ClientProtocol
+    from .datastructures import Headers, HeadersLike, MultipleValuesError
+    from .exceptions import (
+        AbortHandshake,
+        ConnectionClosed,
+        ConnectionClosedError,
+        ConnectionClosedOK,
+        DuplicateParameter,
+        InvalidHandshake,
+        InvalidHeader,
+        InvalidHeaderFormat,
+        InvalidHeaderValue,
+        InvalidMessage,
+        InvalidOrigin,
+        InvalidParameterName,
+        InvalidParameterValue,
+        InvalidState,
+        InvalidStatus,
+        InvalidStatusCode,
+        InvalidUpgrade,
+        InvalidURI,
+        NegotiationError,
+        PayloadTooBig,
+        ProtocolError,
+        RedirectHandshake,
+        SecurityError,
+        WebSocketException,
+        WebSocketProtocolError,
+    )
+    from .legacy.auth import (
+        BasicAuthWebSocketServerProtocol,
+        basic_auth_protocol_factory,
+    )
+    from .legacy.client import WebSocketClientProtocol, connect, unix_connect
+    from .legacy.protocol import WebSocketCommonProtocol, broadcast
+    from .legacy.server import (
+        WebSocketServer,
+        WebSocketServerProtocol,
+        serve,
+        unix_serve,
+    )
+    from .server import ServerProtocol
+    from .typing import (
+        Data,
+        ExtensionName,
+        ExtensionParameter,
+        LoggerLike,
+        Origin,
+        Subprotocol,
+    )
+else:
+    lazy_import(
+        globals(),
+        aliases={
+            # .client
+            "ClientProtocol": ".client",
+            # .datastructures
+            "Headers": ".datastructures",
+            "HeadersLike": ".datastructures",
+            "MultipleValuesError": ".datastructures",
+            # .exceptions
+            "AbortHandshake": ".exceptions",
+            "ConnectionClosed": ".exceptions",
+            "ConnectionClosedError": ".exceptions",
+            "ConnectionClosedOK": ".exceptions",
+            "DuplicateParameter": ".exceptions",
+            "InvalidHandshake": ".exceptions",
+            "InvalidHeader": ".exceptions",
+            "InvalidHeaderFormat": ".exceptions",
+            "InvalidHeaderValue": ".exceptions",
+            "InvalidMessage": ".exceptions",
+            "InvalidOrigin": ".exceptions",
+            "InvalidParameterName": ".exceptions",
+            "InvalidParameterValue": ".exceptions",
+            "InvalidState": ".exceptions",
+            "InvalidStatus": ".exceptions",
+            "InvalidStatusCode": ".exceptions",
+            "InvalidUpgrade": ".exceptions",
+            "InvalidURI": ".exceptions",
+            "NegotiationError": ".exceptions",
+            "PayloadTooBig": ".exceptions",
+            "ProtocolError": ".exceptions",
+            "RedirectHandshake": ".exceptions",
+            "SecurityError": ".exceptions",
+            "WebSocketException": ".exceptions",
+            "WebSocketProtocolError": ".exceptions",
+            # .legacy.auth
+            "BasicAuthWebSocketServerProtocol": ".legacy.auth",
+            "basic_auth_protocol_factory": ".legacy.auth",
+            # .legacy.client
+            "WebSocketClientProtocol": ".legacy.client",
+            "connect": ".legacy.client",
+            "unix_connect": ".legacy.client",
+            # .legacy.protocol
+            "WebSocketCommonProtocol": ".legacy.protocol",
+            "broadcast": ".legacy.protocol",
+            # .legacy.server
+            "WebSocketServer": ".legacy.server",
+            "WebSocketServerProtocol": ".legacy.server",
+            "serve": ".legacy.server",
+            "unix_serve": ".legacy.server",
+            # .server
+            "ServerProtocol": ".server",
+            # .typing
+            "Data": ".typing",
+            "ExtensionName": ".typing",
+            "ExtensionParameter": ".typing",
+            "LoggerLike": ".typing",
+            "Origin": ".typing",
+            "Subprotocol": ".typing",
+        },
+        deprecated_aliases={
+            "framing": ".legacy",
+            "handshake": ".legacy",
+            "parse_uri": ".uri",
+            "WebSocketURI": ".uri",
+        },
+    )
diff --git a/src/websockets/http.py b/src/websockets/http.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import sys
+import typing
 
 from .imports import lazy_import
 from .version import version as websockets_version
@@ -9,18 +10,22 @@
 # For backwards compatibility:
 
 
-lazy_import(
-    globals(),
-    # Headers and MultipleValuesError used to be defined in this module.
-    aliases={
-        "Headers": ".datastructures",
-        "MultipleValuesError": ".datastructures",
-    },
-    deprecated_aliases={
-        "read_request": ".legacy.http",
-        "read_response": ".legacy.http",
-    },
-)
+# When type checking, import non-deprecated aliases eagerly. Else, import on demand.
+if typing.TYPE_CHECKING:
+    from .datastructures import Headers, MultipleValuesError  # noqa: F401
+else:
+    lazy_import(
+        globals(),
+        # Headers and MultipleValuesError used to be defined in this module.
+        aliases={
+            "Headers": ".datastructures",
+            "MultipleValuesError": ".datastructures",
+        },
+        deprecated_aliases={
+            "read_request": ".legacy.http",
+            "read_response": ".legacy.http",
+        },
+    )
 
 
 __all__ = ["USER_AGENT"]

diff --git a/tests/test_exports.py b/tests/test_exports.py
@@ -3,6 +3,7 @@
 import websockets
 import websockets.auth
 import websockets.client
+import websockets.datastructures
 import websockets.exceptions
 import websockets.legacy.protocol
 import websockets.server
@@ -13,11 +14,11 @@
 combined_exports = (
     websockets.auth.__all__
     + websockets.client.__all__
+    + websockets.datastructures.__all__
     + websockets.exceptions.__all__
     + websockets.legacy.protocol.__all__
     + websockets.server.__all__
     + websockets.typing.__all__
-    + websockets.uri.__all__
 )