PYTHON-4590 - Add type guards to async API methods (mongodb#1820)

Author: NoahStapp

pymongo/asynchronous/collection.py

@@ -228,6 +228,10 @@ 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 but {type(database)} given")
 
  if not name or ".." in name:
  raise InvalidName("collection names cannot be empty")

pymongo/asynchronous/database.py

@@ -119,9 +119,14 @@ 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 given {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
@@ -598,6 +596,9 @@ 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 given {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 +684,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 {type(database)} given"
+ )
+
  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,9 @@ class _MongoClientErrorHandler:
  def __init__(
  self, client: AsyncMongoClient, server: Server, session: Optional[AsyncClientSession]
  ):
+ if not isinstance(client, AsyncMongoClient):
+ raise TypeError(f"AsyncMongoClient required but given {type(client)}")
+
  self.client = client
  self.server_address = server.description.address
  self.session = session

pymongo/synchronous/collection.py

@@ -231,6 +231,10 @@ 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 but {type(database)} given")
 
  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 given {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
@@ -596,6 +594,9 @@ 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 given {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 +682,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 {type(database)} given"
+ )
+
  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 given {type(client)}")
+
  self.client = client
  self.server_address = server.description.address
  self.session = session

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: