PYTHON-4260 Lazily load optional imports (#1550)

Author: blink1073

Diff for: .evergreen/config.yml

@@ -2071,6 +2071,21 @@ tasks:
               bash $SCRIPT -p $CONFIG -h ${github_commit} -o "mongodb" -n "mongo-python-driver"
               echo '{"results": [{ "status": "PASS", "test_file": "Build", "log_raw": "Test completed"  } ]}' > ${PROJECT_DIRECTORY}/test-results.json
 
+    - name: "check-import-time"
+      tags: ["pr"]
+      commands:
+        - command: shell.exec
+          type: test
+          params:
+            shell: "bash"
+            working_dir: src
+            script: |
+              ${PREPARE_SHELL}
+              set -x
+              export BASE_SHA=${revision}
+              export HEAD_SHA=${github_commit}
+              bash .evergreen/run-import-time-test.sh
+
 axes:
   # Choice of distro
   - id: platform
@@ -3046,6 +3061,12 @@ buildvariants:
   tasks:
     - name: "assign-pr-reviewer"
 
+- name: rhel8-import-time
+  display_name: Import Time Check
+  run_on: rhel87-small
+  tasks:
+    - name: "check-import-time"
+
 - name: Release
   display_name: Release
   batchtime: 20160 # 14 days

Diff for: .evergreen/run-import-time-test.sh

@@ -0,0 +1,31 @@
+#!/bin/bash -ex
+
+set -o errexit  # Exit the script with error if any of the commands fail
+set -x
+
+. .evergreen/utils.sh
+
+if [ -z "$PYTHON_BINARY" ]; then
+    PYTHON_BINARY=$(find_python3)
+fi
+
+# Use the previous commit if this was not a PR run.
+if [ "$BASE_SHA" == "$HEAD_SHA" ]; then
+    BASE_SHA=$(git rev-parse HEAD~1)
+fi
+
+function get_import_time() {
+    local log_file
+    createvirtualenv "$PYTHON_BINARY" import-venv
+    python -m pip install -q ".[aws,encryption,gssapi,ocsp,snappy,zstd]"
+    # Import once to cache modules
+    python -c "import pymongo"
+    log_file="pymongo-$1.log"
+    python -X importtime -c "import pymongo" 2> $log_file
+}
+
+get_import_time $HEAD_SHA
+git checkout $BASE_SHA
+get_import_time $BASE_SHA
+git checkout $HEAD_SHA
+python tools/compare_import_time.py $HEAD_SHA $BASE_SHA

Diff for: .evergreen/run-tests.sh

@@ -248,7 +248,9 @@ if [ -n "$COVERAGE" ] && [ "$PYTHON_IMPL" = "CPython" ]; then
 fi
 
 if [ -n "$GREEN_FRAMEWORK" ]; then
-     python -m pip install $GREEN_FRAMEWORK
+    # Install all optional deps to ensure lazy imports are getting patched.
+    python -m pip install -q ".[aws,encryption,gssapi,ocsp,snappy,zstd]"
+    python -m pip install $GREEN_FRAMEWORK
 fi
 
 # Show the installed packages

Diff for: doc/changelog.rst

@@ -74,6 +74,12 @@ Unavoidable breaking changes
     >>> dict_to_SON(data_as_dict)
     SON([('driver', SON([('name', 'PyMongo'), ('version', '4.7.0.dev0')])), ('os', SON([('type', 'Darwin'), ('name', 'Darwin'), ('architecture', 'arm64'), ('version', '14.3')])), ('platform', 'CPython 3.11.6.final.0')])
 
+- PyMongo now uses `lazy imports <https://docs.python.org/3/library/importlib.html#implementing-lazy-imports>`_ for external dependencies.
+  If you are relying on any kind of monkey-patching of the standard library, you may need to explicitly import those external libraries in addition
+  to ``pymongo`` before applying the patch.  Note that we test with ``gevent`` and ``eventlet`` patching, and those continue to work.
+
+- The "aws" extra now requires minimum version of ``1.1.0`` for ``pymongo_auth_aws``.
+
 Changes in Version 4.6.2
 ------------------------
 

Diff for: pymongo/_azure_helpers.py

@@ -17,12 +17,14 @@
 
 import json
 from typing import Any, Optional
-from urllib.request import Request, urlopen
 
 
 def _get_azure_response(
     resource: str, client_id: Optional[str] = None, timeout: float = 5
 ) -> dict[str, Any]:
+    # Deferred import to save overall import time.
+    from urllib.request import Request, urlopen
+
     url = "http://169.254.169.254/metadata/identity/oauth2/token"
     url += "?api-version=2018-02-01"
     url += f"&resource={resource}"

Diff for: pymongo/_csot.py

@@ -21,9 +21,10 @@
 from collections import deque
 from contextlib import AbstractContextManager
 from contextvars import ContextVar, Token
-from typing import Any, Callable, Deque, MutableMapping, Optional, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, Deque, MutableMapping, Optional, TypeVar, cast
 
-from pymongo.write_concern import WriteConcern
+if TYPE_CHECKING:
+    from pymongo.write_concern import WriteConcern
 
 TIMEOUT: ContextVar[Optional[float]] = ContextVar("TIMEOUT", default=None)
 RTT: ContextVar[float] = ContextVar("RTT", default=0.0)

Diff for: pymongo/_lazy_import.py

@@ -0,0 +1,38 @@
+# Copyright 2024-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you
+# may not use this file except in compliance with the License.  You
+# may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.  See the License for the specific language governing
+# permissions and limitations under the License.
+from __future__ import annotations
+
+import importlib.util
+import sys
+from types import ModuleType
+
+
+def lazy_import(name: str) -> ModuleType:
+    """Lazily import a module by name
+
+    From https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
+    """
+    try:
+        spec = importlib.util.find_spec(name)
+    except ValueError:
+        raise ModuleNotFoundError(name=name) from None
+    if spec is None:
+        raise ModuleNotFoundError(name=name)
+    assert spec is not None
+    loader = importlib.util.LazyLoader(spec.loader)  # type:ignore[arg-type]
+    spec.loader = loader
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    loader.exec_module(module)
+    return module

Diff for: pymongo/auth_aws.py

@@ -15,38 +15,16 @@
 """MONGODB-AWS Authentication helpers."""
 from __future__ import annotations
 
-try:
-    import pymongo_auth_aws  # type:ignore[import]
-    from pymongo_auth_aws import (
-        AwsCredential,
-        AwsSaslContext,
-        PyMongoAuthAwsError,
-    )
+from pymongo._lazy_import import lazy_import
 
+try:
+    pymongo_auth_aws = lazy_import("pymongo_auth_aws")
     _HAVE_MONGODB_AWS = True
 except ImportError:
-
-    class AwsSaslContext:  # type: ignore
-        def __init__(self, credentials: MongoCredential):
-            pass
-
     _HAVE_MONGODB_AWS = False
 
-try:
-    from pymongo_auth_aws.auth import (  # type:ignore[import]
-        set_cached_credentials,
-        set_use_cached_credentials,
-    )
-
-    # Enable credential caching.
-    set_use_cached_credentials(True)
-except ImportError:
-
-    def set_cached_credentials(_creds: Optional[AwsCredential]) -> None:
-        pass
 
-
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Type
+from typing import TYPE_CHECKING, Any, Mapping, Type
 
 import bson
 from bson.binary import Binary
@@ -58,21 +36,6 @@ def set_cached_credentials(_creds: Optional[AwsCredential]) -> None:
     from pymongo.pool import Connection
 
 
-class _AwsSaslContext(AwsSaslContext):  # type: ignore
-    # Dependency injection:
-    def binary_type(self) -> Type[Binary]:
-        """Return the bson.binary.Binary type."""
-        return Binary
-
-    def bson_encode(self, doc: Mapping[str, Any]) -> bytes:
-        """Encode a dictionary to BSON."""
-        return bson.encode(doc)
-
-    def bson_decode(self, data: _ReadableBuffer) -> Mapping[str, Any]:
-        """Decode BSON to a dictionary."""
-        return bson.decode(data)
-
-
 def _authenticate_aws(credentials: MongoCredential, conn: Connection) -> None:
     """Authenticate using MONGODB-AWS."""
     if not _HAVE_MONGODB_AWS:
@@ -84,9 +47,23 @@ def _authenticate_aws(credentials: MongoCredential, conn: Connection) -> None:
     if conn.max_wire_version < 9:
         raise ConfigurationError("MONGODB-AWS authentication requires MongoDB version 4.4 or later")
 
+    class AwsSaslContext(pymongo_auth_aws.AwsSaslContext):  # type: ignore
+        # Dependency injection:
+        def binary_type(self) -> Type[Binary]:
+            """Return the bson.binary.Binary type."""
+            return Binary
+
+        def bson_encode(self, doc: Mapping[str, Any]) -> bytes:
+            """Encode a dictionary to BSON."""
+            return bson.encode(doc)
+
+        def bson_decode(self, data: _ReadableBuffer) -> Mapping[str, Any]:
+            """Decode BSON to a dictionary."""
+            return bson.decode(data)
+
     try:
-        ctx = _AwsSaslContext(
-            AwsCredential(
+        ctx = AwsSaslContext(
+            pymongo_auth_aws.AwsCredential(
                 credentials.username,
                 credentials.password,
                 credentials.mechanism_properties.aws_session_token,
@@ -108,14 +85,14 @@ def _authenticate_aws(credentials: MongoCredential, conn: Connection) -> None:
             if res["done"]:
                 # SASL complete.
                 break
-    except PyMongoAuthAwsError as exc:
+    except pymongo_auth_aws.PyMongoAuthAwsError as exc:
         # Clear the cached credentials if we hit a failure in auth.
-        set_cached_credentials(None)
+        pymongo_auth_aws.set_cached_credentials(None)
         # Convert to OperationFailure and include pymongo-auth-aws version.
         raise OperationFailure(
             f"{exc} (pymongo-auth-aws version {pymongo_auth_aws.__version__})"
         ) from None
     except Exception:
         # Clear the cached credentials if we hit a failure in auth.
-        set_cached_credentials(None)
+        pymongo_auth_aws.set_cached_credentials(None)
         raise

Diff for: pymongo/compression_support.py

@@ -16,32 +16,31 @@
 import warnings
 from typing import Any, Iterable, Optional, Union
 
-try:
-    import snappy  # type:ignore[import]
+from pymongo._lazy_import import lazy_import
+from pymongo.hello import HelloCompat
+from pymongo.monitoring import _SENSITIVE_COMMANDS
 
+try:
+    snappy = lazy_import("snappy")
     _HAVE_SNAPPY = True
 except ImportError:
     # python-snappy isn't available.
     _HAVE_SNAPPY = False
 
 try:
-    import zlib
+    zlib = lazy_import("zlib")
 
     _HAVE_ZLIB = True
 except ImportError:
     # Python built without zlib support.
     _HAVE_ZLIB = False
 
 try:
-    from zstandard import ZstdCompressor, ZstdDecompressor
-
+    zstandard = lazy_import("zstandard")
     _HAVE_ZSTD = True
 except ImportError:
     _HAVE_ZSTD = False
 
-from pymongo.hello import HelloCompat
-from pymongo.monitoring import _SENSITIVE_COMMANDS
-
 _SUPPORTED_COMPRESSORS = {"snappy", "zlib", "zstd"}
 _NO_COMPRESSION = {HelloCompat.CMD, HelloCompat.LEGACY_CMD}
 _NO_COMPRESSION.update(_SENSITIVE_COMMANDS)
@@ -138,7 +137,7 @@ class ZstdContext:
     def compress(data: bytes) -> bytes:
         # ZstdCompressor is not thread safe.
         # TODO: Use a pool?
-        return ZstdCompressor().compress(data)
+        return zstandard.ZstdCompressor().compress(data)
 
 
 def decompress(data: bytes, compressor_id: int) -> bytes:
@@ -153,6 +152,6 @@ def decompress(data: bytes, compressor_id: int) -> bytes:
     elif compressor_id == ZstdContext.compressor_id:
         # ZstdDecompressor is not thread safe.
         # TODO: Use a pool?
-        return ZstdDecompressor().decompress(data)
+        return zstandard.ZstdDecompressor().decompress(data)
     else:
         raise ValueError("Unknown compressorId %d" % (compressor_id,))

Diff for: pymongo/errors.py

@@ -15,24 +15,14 @@
 """Exceptions raised by PyMongo."""
 from __future__ import annotations
 
+from ssl import SSLCertVerificationError as _CertificateError  # noqa: F401
 from typing import TYPE_CHECKING, Any, Iterable, Mapping, Optional, Sequence, Union
 
 from bson.errors import InvalidDocument
 
 if TYPE_CHECKING:
     from pymongo.typings import _DocumentOut
 
-try:
-    # CPython 3.7+
-    from ssl import SSLCertVerificationError as _CertificateError
-except ImportError:
-    try:
-        from ssl import CertificateError as _CertificateError
-    except ImportError:
-
-        class _CertificateError(ValueError):  # type: ignore
-            pass
-
 
 class PyMongoError(Exception):
     """Base class for all PyMongo exceptions."""