Add helper to manage aliases and deprecations.

This may save a little bit of CPU and memory by avoiding
unnecessary imports too, especially as the library grows.

Author: aaugustin

src/websockets/__init__.py

@@ -1,19 +1,8 @@
-# This relies on each of the submodules having an __all__ variable.
-
-from .client import *
-from .datastructures import *  # noqa
-from .exceptions import *  # noqa
-from .legacy.auth import *  # noqa
-from .legacy.client import *  # noqa
-from .legacy.protocol import *  # noqa
-from .legacy.server import *  # noqa
-from .server import *
-from .typing import *  # noqa
-from .uri import *  # noqa
+from .imports import lazy_import
 from .version import version as __version__  # noqa
 
 
-__all__ = [
+__all__ = [  # noqa
     "AbortHandshake",
     "basic_auth_protocol_factory",
     "BasicAuthWebSocketServerProtocol",
@@ -58,3 +47,60 @@
     "WebSocketServerProtocol",
     "WebSocketURI",
 ]
+
+lazy_import(
+    globals(),
+    aliases={
+        "auth": ".legacy",
+        "basic_auth_protocol_factory": ".legacy.auth",
+        "BasicAuthWebSocketServerProtocol": ".legacy.auth",
+        "ClientConnection": ".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",
+        "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",
+        "ServerConnection": ".server",
+        "serve": ".legacy.server",
+        "unix_serve": ".legacy.server",
+        "WebSocketServerProtocol": ".legacy.server",
+        "WebSocketServer": ".legacy.server",
+        "Data": ".typing",
+        "Origin": ".typing",
+        "ExtensionHeader": ".typing",
+        "ExtensionParameter": ".typing",
+        "Subprotocol": ".typing",
+        "parse_uri": ".uri",
+        "WebSocketURI": ".uri",
+    },
+    deprecated_aliases={
+        "framing": ".legacy",
+        "handshake": ".legacy",
+    },
+)

src/websockets/auth.py

@@ -24,7 +24,7 @@
 )
 from .http import USER_AGENT, build_host
 from .http11 import Request, Response
-from .legacy.client import WebSocketClientProtocol, connect, unix_connect  # noqa
+from .imports import lazy_import
 from .typing import (
     ConnectionOption,
     ExtensionHeader,
@@ -36,6 +36,15 @@
 from .utils import accept_key, generate_key
 
 
+lazy_import(
+    globals(),
+    aliases={
+        "connect": ".legacy.client",
+        "unix_connect": ".legacy.client",
+        "WebSocketClientProtocol": ".legacy.client",
+    },
+)
+
 __all__ = ["ClientConnection"]
 
 logger = logging.getLogger(__name__)

src/websockets/client.py

@@ -1,15 +1,27 @@
-import asyncio
 import ipaddress
 import sys
-import warnings
-from typing import Tuple
 
-# For backwards compatibility:
-# Headers and MultipleValuesError used to be defined in this module
-from .datastructures import Headers, MultipleValuesError  # noqa
+from .imports import lazy_import
 from .version import version as websockets_version
 
 
+# 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",
+    },
+)
+
+
 __all__ = ["USER_AGENT", "build_host"]
 
 
@@ -38,24 +50,3 @@ def build_host(host: str, port: int, secure: bool) -> str:
         host = f"{host}:{port}"
 
     return host
-
-
-# Backwards compatibility with previously documented public APIs
-
-
-async def read_request(
-    stream: asyncio.StreamReader,
-) -> Tuple[str, Headers]:  # pragma: no cover
-    warnings.warn("websockets.http.read_request is deprecated", DeprecationWarning)
-    from .legacy.http import read_request
-
-    return await read_request(stream)
-
-
-async def read_response(
-    stream: asyncio.StreamReader,
-) -> Tuple[int, str, Headers]:  # pragma: no cover
-    warnings.warn("websockets.http.read_response is deprecated", DeprecationWarning)
-    from .legacy.http import read_response
-
-    return await read_response(stream)

src/websockets/framing.py

@@ -0,0 +1,95 @@
+import importlib
+import sys
+import warnings
+from typing import Any, Dict, Iterable, Optional
+
+
+__all__ = ["lazy_import"]
+
+
+def lazy_import(
+    namespace: Dict[str, Any],
+    aliases: Optional[Dict[str, str]] = None,
+    deprecated_aliases: Optional[Dict[str, str]] = None,
+) -> None:
+    """
+    Provide lazy, module-level imports.
+
+    Typical use::
+
+        __getattr__, __dir__ = lazy_import(
+            globals(),
+            aliases={
+                "<name>": "<source module>",
+                ...
+            },
+            deprecated_aliases={
+                ...,
+            }
+        )
+
+    This function defines __getattr__ and __dir__ per PEP 562.
+
+    On Python 3.6 and earlier, it falls back to non-lazy imports and doesn't
+    raise deprecation warnings.
+
+    """
+    if aliases is None:
+        aliases = {}
+    if deprecated_aliases is None:
+        deprecated_aliases = {}
+
+    namespace_set = set(namespace)
+    aliases_set = set(aliases)
+    deprecated_aliases_set = set(deprecated_aliases)
+
+    assert not namespace_set & aliases_set, "namespace conflict"
+    assert not namespace_set & deprecated_aliases_set, "namespace conflict"
+    assert not aliases_set & deprecated_aliases_set, "namespace conflict"
+
+    package = namespace["__name__"]
+
+    if sys.version_info[:2] >= (3, 7):
+
+        def __getattr__(name: str) -> Any:
+            assert aliases is not None  # mypy cannot figure this out
+            try:
+                source = aliases[name]
+            except KeyError:
+                pass
+            else:
+                module = importlib.import_module(source, package)
+                return getattr(module, name)
+
+            assert deprecated_aliases is not None  # mypy cannot figure this out
+            try:
+                source = deprecated_aliases[name]
+            except KeyError:
+                pass
+            else:
+                warnings.warn(
+                    f"{package}.{name} is deprecated",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                module = importlib.import_module(source, package)
+                return getattr(module, name)
+
+            raise AttributeError(f"module {package!r} has no attribute {name!r}")
+
+        namespace["__getattr__"] = __getattr__
+
+        def __dir__() -> Iterable[str]:
+            return sorted(namespace_set | aliases_set | deprecated_aliases_set)
+
+        namespace["__dir__"] = __dir__
+
+    else:  # pragma: no cover
+
+        for name, source in aliases.items():
+            module = importlib.import_module(source, package)
+            namespace[name] = getattr(module, name)
+
+        for name, source in deprecated_aliases.items():
+            module = importlib.import_module(source, package)
+            namespace[name] = getattr(module, name)

src/websockets/handshake.py

@@ -26,12 +26,7 @@
 )
 from .http import USER_AGENT
 from .http11 import Request, Response
-from .legacy.server import (  # noqa
-    WebSocketServer,
-    WebSocketServerProtocol,
-    serve,
-    unix_serve,
-)
+from .imports import lazy_import
 from .typing import (
     ConnectionOption,
     ExtensionHeader,
@@ -42,6 +37,17 @@
 from .utils import accept_key
 
 
+lazy_import(
+    globals(),
+    aliases={
+        "serve": ".legacy.server",
+        "unix_serve": ".legacy.server",
+        "WebSocketServerProtocol": ".legacy.server",
+        "WebSocketServer": ".legacy.server",
+    },
+)
+
+
 __all__ = ["ServerConnection"]
 
 logger = logging.getLogger(__name__)