PYTHON-4590 - Add type guards to async API methods

pymongo/asynchronous/bulk.py

@@ -240,6 +240,7 @@ async def write_command(
         client: AsyncMongoClient,
     ) -> dict[str, Any]:
         """A proxy for SocketInfo.write_command that handles event publishing."""
+
         cmd[bwc.field] = docs
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
             _debug_log(
@@ -324,6 +325,7 @@ async def unack_write(
         client: AsyncMongoClient,
     ) -> Optional[Mapping[str, Any]]:
         """A proxy for AsyncConnection.unack_write that handles event publishing."""
+
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
             _debug_log(
                 _COMMAND_LOGGER,

pymongo/asynchronous/collection.py

@@ -228,6 +228,12 @@ def __init__(
         )
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
+        from pymongo.asynchronous.database import AsyncDatabase
+
+        if not isinstance(database, AsyncDatabase):
+            raise TypeError(
+                f"AsyncCollection requires an AsyncDatabase, {database} is an instance of {type(database)}"
+            )
 
         if not name or ".." in name:
             raise InvalidName("collection names cannot be empty")

pymongo/asynchronous/database.py

@@ -119,9 +119,16 @@ def __init__(
             read_concern or client.read_concern,
         )
 
+        from pymongo.asynchronous.mongo_client import AsyncMongoClient
+
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
 
+        if not isinstance(client, AsyncMongoClient):
+            raise TypeError(
+                f"AsyncMongoClient required but {client} is an instance of {type(client)}"
+            )
+
         if name != "$external":
             _check_name(name)
 

pymongo/asynchronous/encryption.py

@@ -194,9 +194,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             # Wrap I/O errors in PyMongo exceptions.
             _raise_connection_failure((host, port), error)
 
-    async def collection_info(
-        self, database: AsyncDatabase[Mapping[str, Any]], filter: bytes
-    ) -> Optional[bytes]:
+    async def collection_info(self, database: str, filter: bytes) -> Optional[bytes]:
         """Get the collection info for a namespace.
 
         The returned collection info is passed to libmongocrypt which reads
@@ -339,6 +337,7 @@ def __init__(self, client: AsyncMongoClient[_DocumentTypeArg], opts: AutoEncrypt
         :param client: The encrypted AsyncMongoClient.
         :param opts: The encrypted client's :class:`AutoEncryptionOpts`.
         """
+
         if opts._schema_map is None:
             schema_map = None
         else:
@@ -598,6 +597,11 @@ def __init__(
         if not isinstance(codec_options, CodecOptions):
             raise TypeError("codec_options must be an instance of bson.codec_options.CodecOptions")
 
+        if not isinstance(key_vault_client, AsyncMongoClient):
+            raise TypeError(
+                f"AsyncMongoClient required but {key_vault_client} is an instance of {type(key_vault_client)}"
+            )
+
         self._kms_providers = kms_providers
         self._key_vault_namespace = key_vault_namespace
         self._key_vault_client = key_vault_client
@@ -683,6 +687,11 @@ async def create_encrypted_collection(
             https://mongodb.com/docs/manual/reference/command/create
 
         """
+        if not isinstance(database, AsyncDatabase):
+            raise TypeError(
+                f"create_encrypted_collection() requires an AsyncDatabase, but {database} is an instance of {type(database)}"
+            )
+
         encrypted_fields = deepcopy(encrypted_fields)
         for i, field in enumerate(encrypted_fields["fields"]):
             if isinstance(field, dict) and field.get("keyId") is None:

pymongo/asynchronous/mongo_client.py

@@ -2419,6 +2419,11 @@ class _MongoClientErrorHandler:
     def __init__(
         self, client: AsyncMongoClient, server: Server, session: Optional[AsyncClientSession]
     ):
+        if not isinstance(client, AsyncMongoClient):
+            raise TypeError(
+                f"AsyncMongoClient required but {client} is an instance of {type(client)}"
+            )
+
         self.client = client
         self.server_address = server.description.address
         self.session = session

pymongo/asynchronous/server.py

@@ -157,7 +157,9 @@ async def run_operation(
         :param unpack_res: A callable that decodes the wire protocol response.
         :param client: An AsyncMongoClient instance.
         """
+
         assert listeners is not None
+
         publish = listeners.enabled_for_commands
         start = datetime.now()
 

pymongo/synchronous/bulk.py

@@ -240,6 +240,7 @@ def write_command(
         client: MongoClient,
     ) -> dict[str, Any]:
         """A proxy for SocketInfo.write_command that handles event publishing."""
+
         cmd[bwc.field] = docs
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
             _debug_log(
@@ -324,6 +325,7 @@ def unack_write(
         client: MongoClient,
     ) -> Optional[Mapping[str, Any]]:
         """A proxy for Connection.unack_write that handles event publishing."""
+
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
             _debug_log(
                 _COMMAND_LOGGER,

pymongo/synchronous/collection.py

@@ -231,6 +231,12 @@ def __init__(
         )
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
+        from pymongo.synchronous.database import Database
+
+        if not isinstance(database, Database):
+            raise TypeError(
+                f"Collection requires a Database, {database} is an instance of {type(database)}"
+            )
 
         if not name or ".." in name:
             raise InvalidName("collection names cannot be empty")

pymongo/synchronous/database.py

@@ -119,9 +119,14 @@ def __init__(
             read_concern or client.read_concern,
         )
 
+        from pymongo.synchronous.mongo_client import MongoClient
+
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
 
+        if not isinstance(client, MongoClient):
+            raise TypeError(f"MongoClient required but {client} is an instance of {type(client)}")
+
         if name != "$external":
             _check_name(name)
 

pymongo/synchronous/encryption.py

@@ -194,9 +194,7 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             # Wrap I/O errors in PyMongo exceptions.
             _raise_connection_failure((host, port), error)
 
-    def collection_info(
-        self, database: Database[Mapping[str, Any]], filter: bytes
-    ) -> Optional[bytes]:
+    def collection_info(self, database: str, filter: bytes) -> Optional[bytes]:
         """Get the collection info for a namespace.
 
         The returned collection info is passed to libmongocrypt which reads
@@ -337,6 +335,7 @@ def __init__(self, client: MongoClient[_DocumentTypeArg], opts: AutoEncryptionOp
         :param client: The encrypted MongoClient.
         :param opts: The encrypted client's :class:`AutoEncryptionOpts`.
         """
+
         if opts._schema_map is None:
             schema_map = None
         else:
@@ -596,6 +595,11 @@ def __init__(
         if not isinstance(codec_options, CodecOptions):
             raise TypeError("codec_options must be an instance of bson.codec_options.CodecOptions")
 
+        if not isinstance(key_vault_client, MongoClient):
+            raise TypeError(
+                f"MongoClient required but {key_vault_client} is an instance of {type(key_vault_client)}"
+            )
+
         self._kms_providers = kms_providers
         self._key_vault_namespace = key_vault_namespace
         self._key_vault_client = key_vault_client
@@ -681,6 +685,11 @@ def create_encrypted_collection(
             https://mongodb.com/docs/manual/reference/command/create
 
         """
+        if not isinstance(database, Database):
+            raise TypeError(
+                f"create_encrypted_collection() requires a Database, but {database} is an instance of {type(database)}"
+            )
+
         encrypted_fields = deepcopy(encrypted_fields)
         for i, field in enumerate(encrypted_fields["fields"]):
             if isinstance(field, dict) and field.get("keyId") is None:

pymongo/synchronous/mongo_client.py

@@ -2406,6 +2406,9 @@ class _MongoClientErrorHandler:
     )
 
     def __init__(self, client: MongoClient, server: Server, session: Optional[ClientSession]):
+        if not isinstance(client, MongoClient):
+            raise TypeError(f"MongoClient required but {client} is an instance of {type(client)}")
+
         self.client = client
         self.server_address = server.description.address
         self.session = session

pymongo/synchronous/server.py

@@ -157,7 +157,9 @@ def run_operation(
         :param unpack_res: A callable that decodes the wire protocol response.
         :param client: A MongoClient instance.
         """
+
         assert listeners is not None
+
         publish = listeners.enabled_for_commands
         start = datetime.now()
 

test/helpers.py

@@ -35,23 +35,13 @@
     HAVE_IPADDRESS = True
 except ImportError:
     HAVE_IPADDRESS = False
-from contextlib import contextmanager
 from functools import wraps
-from test.version import Version
 from typing import Any, Callable, Dict, Generator, no_type_check
 from unittest import SkipTest
-from urllib.parse import quote_plus
 
-import pymongo
-import pymongo.errors
 from bson.son import SON
 from pymongo import common, message
-from pymongo.common import partition_node
-from pymongo.hello import HelloCompat
-from pymongo.server_api import ServerApi
 from pymongo.ssl_support import HAVE_SSL, _ssl  # type:ignore[attr-defined]
-from pymongo.synchronous.database import Database
-from pymongo.synchronous.mongo_client import MongoClient
 from pymongo.uri_parser import parse_uri
 
 if HAVE_SSL: