PYTHON-3454 Specifying a generic type for a collection does not correctly enforce type safety when inserting data (#1081)

Author: juliusgeo

.github/workflows/test-python.yml

@@ -67,6 +67,8 @@ jobs:
           # Test overshadowed codec_options.py file
           mypy --install-types --non-interactive bson/codec_options.py
           mypy --install-types --non-interactive --disable-error-code var-annotated --disable-error-code attr-defined --disable-error-code union-attr --disable-error-code assignment --disable-error-code no-redef --disable-error-code index --allow-redefinition --allow-untyped-globals --exclude "test/mypy_fails/*.*" test
+          python -m pip install -U typing_extensions
+          mypy --install-types --non-interactive test/test_mypy.py
 
   linkcheck:
     name: Check Links

doc/examples/type_hints.rst

@@ -92,7 +92,9 @@ Note that when using :class:`~bson.son.SON`, the key and value types must be giv
 Typed Collection
 ----------------
 
-You can use :py:class:`~typing.TypedDict` (Python 3.8+) when using a well-defined schema for the data in a :class:`~pymongo.collection.Collection`:
+You can use :py:class:`~typing.TypedDict` (Python 3.8+) when using a well-defined schema for the data in a
+:class:`~pymongo.collection.Collection`. Note that all `schema validation`_ for inserts and updates is done on the server.
+These methods automatically add an "_id" field.
 
 .. doctest::
 
@@ -105,10 +107,12 @@ You can use :py:class:`~typing.TypedDict` (Python 3.8+) when using a well-define
   ...
   >>> client: MongoClient = MongoClient()
   >>> collection: Collection[Movie] = client.test.test
-  >>> inserted = collection.insert_one({"name": "Jurassic Park", "year": 1993 })
+  >>> inserted = collection.insert_one(Movie(name="Jurassic Park", year=1993))
   >>> result = collection.find_one({"name": "Jurassic Park"})
   >>> assert result is not None
   >>> assert result["year"] == 1993
+  >>> # This will not be type checked, despite being present, because it is added by PyMongo.
+  >>> assert type(result["_id"]) == ObjectId
 
 Typed Database
 --------------
@@ -243,3 +247,4 @@ Another example is trying to set a value on a :class:`~bson.raw_bson.RawBSONDocu
 .. _limitations in mypy: https://github.com/python/mypy/issues/3737
 .. _mypy config: https://mypy.readthedocs.io/en/stable/config_file.html
 .. _test_mypy module: https://github.com/mongodb/mongo-python-driver/blob/master/test/test_mypy.py
+.. _schema validation: https://www.mongodb.com/docs/manual/core/schema-validation/#when-to-use-schema-validation

pymongo/client_session.py

@@ -435,7 +435,7 @@ def _max_time_expired_error(exc):
 
 # From the transactions spec, all the retryable writes errors plus
 # WriteConcernFailed.
-_UNKNOWN_COMMIT_ERROR_CODES = _RETRYABLE_ERROR_CODES | frozenset(
+_UNKNOWN_COMMIT_ERROR_CODES: frozenset = _RETRYABLE_ERROR_CODES | frozenset(
     [
         64,  # WriteConcernFailed
         50,  # MaxTimeMSExpired

pymongo/collection.py

@@ -71,7 +71,7 @@
     InsertOneResult,
     UpdateResult,
 )
-from pymongo.typings import _CollationIn, _DocumentIn, _DocumentType, _Pipeline
+from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
 from pymongo.write_concern import WriteConcern
 
 _FIND_AND_MODIFY_DOC_FIELDS = {"value": 1}
@@ -566,7 +566,7 @@ def _insert_command(session, sock_info, retryable_write):
 
     def insert_one(
         self,
-        document: _DocumentIn,
+        document: Union[_DocumentType, RawBSONDocument],
         bypass_document_validation: bool = False,
         session: Optional["ClientSession"] = None,
         comment: Optional[Any] = None,
@@ -614,7 +614,7 @@ def insert_one(
         """
         common.validate_is_document_type("document", document)
         if not (isinstance(document, RawBSONDocument) or "_id" in document):
-            document["_id"] = ObjectId()
+            document["_id"] = ObjectId()  # type: ignore[index]
 
         write_concern = self._write_concern_for(session)
         return InsertOneResult(
@@ -633,7 +633,7 @@ def insert_one(
     @_csot.apply
     def insert_many(
         self,
-        documents: Iterable[_DocumentIn],
+        documents: Iterable[Union[_DocumentType, RawBSONDocument]],
         ordered: bool = True,
         bypass_document_validation: bool = False,
         session: Optional["ClientSession"] = None,
@@ -697,7 +697,7 @@ def gen():
                 common.validate_is_document_type("document", document)
                 if not isinstance(document, RawBSONDocument):
                     if "_id" not in document:
-                        document["_id"] = ObjectId()
+                        document["_id"] = ObjectId()  # type: ignore[index]
                     inserted_ids.append(document["_id"])
                 yield (message._INSERT, document)
 

pymongo/helpers.py

@@ -44,7 +44,7 @@
 # From the SDAM spec, the "not primary" error codes are combined with the
 # "node is recovering" error codes (of which the "node is shutting down"
 # errors are a subset).
-_NOT_PRIMARY_CODES = (
+_NOT_PRIMARY_CODES: frozenset = (
     frozenset(
         [
             10058,  # LegacyNotPrimary <=3.2 "not primary" error code
@@ -58,7 +58,7 @@
     | _SHUTDOWN_CODES
 )
 # From the retryable writes spec.
-_RETRYABLE_ERROR_CODES = _NOT_PRIMARY_CODES | frozenset(
+_RETRYABLE_ERROR_CODES: frozenset = _NOT_PRIMARY_CODES | frozenset(
     [
         7,  # HostNotFound
         6,  # HostUnreachable

pymongo/monitoring.py

@@ -528,7 +528,7 @@ def register(listener: _EventListener) -> None:
 # Note - to avoid bugs from forgetting which if these is all lowercase and
 # which are camelCase, and at the same time avoid having to add a test for
 # every command, use all lowercase here and test against command_name.lower().
-_SENSITIVE_COMMANDS = set(
+_SENSITIVE_COMMANDS: set = set(
     [
         "authenticate",
         "saslstart",

test/__init__.py

@@ -43,11 +43,10 @@
     HAVE_IPADDRESS = True
 except ImportError:
     HAVE_IPADDRESS = False
-
 from contextlib import contextmanager
 from functools import wraps
 from test.version import Version
-from typing import Callable, Dict, Generator, no_type_check
+from typing import Any, Callable, Dict, Generator, no_type_check
 from unittest import SkipTest
 from urllib.parse import quote_plus
 
@@ -331,7 +330,9 @@ def hello(self):
 
     def _connect(self, host, port, **kwargs):
         kwargs.update(self.default_client_options)
-        client = pymongo.MongoClient(host, port, serverSelectionTimeoutMS=5000, **kwargs)
+        client: MongoClient = pymongo.MongoClient(
+            host, port, serverSelectionTimeoutMS=5000, **kwargs
+        )
         try:
             try:
                 client.admin.command(HelloCompat.LEGACY_CMD)  # Can we connect?
@@ -356,7 +357,7 @@ def _init_client(self):
         if self.client is not None:
             # Return early when connected to dataLake as mongohoused does not
             # support the getCmdLineOpts command and is tested without TLS.
-            build_info = self.client.admin.command("buildInfo")
+            build_info: Any = self.client.admin.command("buildInfo")
             if "dataLake" in build_info:
                 self.is_data_lake = True
                 self.auth_enabled = True
@@ -521,14 +522,16 @@ def has_secondaries(self):
     @property
     def storage_engine(self):
         try:
-            return self.server_status.get("storageEngine", {}).get("name")
+            return self.server_status.get("storageEngine", {}).get(  # type:ignore[union-attr]
+                "name"
+            )
         except AttributeError:
             # Raised if self.server_status is None.
             return None
 
     def _check_user_provided(self):
         """Return True if db_user/db_password is already an admin user."""
-        client = pymongo.MongoClient(
+        client: MongoClient = pymongo.MongoClient(
             host,
             port,
             username=db_user,
@@ -694,7 +697,7 @@ def supports_secondary_read_pref(self):
         if self.has_secondaries:
             return True
         if self.is_mongos:
-            shard = self.client.config.shards.find_one()["host"]
+            shard = self.client.config.shards.find_one()["host"]  # type:ignore[index]
             num_members = shard.count(",") + 1
             return num_members > 1
         return False
@@ -1015,12 +1018,12 @@ def fork(
         """
 
         def _print_threads(*args: object) -> None:
-            if _print_threads.called:
+            if _print_threads.called:  # type:ignore[attr-defined]
                 return
-            _print_threads.called = True
+            _print_threads.called = True  # type:ignore[attr-defined]
             print_thread_tracebacks()
 
-        _print_threads.called = False
+        _print_threads.called = False  # type:ignore[attr-defined]
 
         def _target() -> None:
             signal.signal(signal.SIGUSR1, _print_threads)

test/test_collection.py

@@ -785,7 +785,7 @@ def test_insert_many_invalid(self):
             db.test.insert_many(1)  # type: ignore[arg-type]
 
         with self.assertRaisesRegex(TypeError, "documents must be a non-empty list"):
-            db.test.insert_many(RawBSONDocument(encode({"_id": 2})))  # type: ignore[arg-type]
+            db.test.insert_many(RawBSONDocument(encode({"_id": 2})))
 
     def test_delete_one(self):
         self.db.test.drop()

test/test_mypy.py

@@ -14,22 +14,20 @@
 
 """Test that each file in mypy_fails/ actually fails mypy, and test some
 sample client code that uses PyMongo typings."""
-
 import os
 import tempfile
 import unittest
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List
 
 try:
-    from typing import TypedDict  # type: ignore[attr-defined]
+    from typing_extensions import TypedDict
 
-    # Not available in Python 3.7
     class Movie(TypedDict):  # type: ignore[misc]
         name: str
         year: int
 
 except ImportError:
-    TypeDict = None
+    TypedDict = None
 
 
 try:
@@ -304,6 +302,28 @@ def test_typeddict_document_type(self) -> None:
         assert retreived["year"] == 1
         assert retreived["name"] == "a"
 
+    @only_type_check
+    def test_typeddict_document_type_insertion(self) -> None:
+        client: MongoClient[Movie] = MongoClient()
+        coll = client.test.test
+        mov = {"name": "THX-1138", "year": 1971}
+        movie = Movie(name="THX-1138", year=1971)
+        coll.insert_one(mov)  # type: ignore[arg-type]
+        coll.insert_one({"name": "THX-1138", "year": 1971})  # This will work because it is in-line.
+        coll.insert_one(movie)
+        coll.insert_many([mov])  # type: ignore[list-item]
+        coll.insert_many([movie])
+        bad_mov = {"name": "THX-1138", "year": "WRONG TYPE"}
+        bad_movie = Movie(name="THX-1138", year="WRONG TYPE")  # type: ignore[typeddict-item]
+        coll.insert_one(bad_mov)  # type:ignore[arg-type]
+        coll.insert_one({"name": "THX-1138", "year": "WRONG TYPE"})  # type: ignore[typeddict-item]
+        coll.insert_one(bad_movie)
+        coll.insert_many([bad_mov])  # type: ignore[list-item]
+        coll.insert_many(
+            [{"name": "THX-1138", "year": "WRONG TYPE"}]  # type: ignore[typeddict-item]
+        )
+        coll.insert_many([bad_movie])
+
     @only_type_check
     def test_raw_bson_document_type(self) -> None:
         client = MongoClient(document_class=RawBSONDocument)

test/utils.py

@@ -601,7 +601,7 @@ def ensure_all_connected(client: MongoClient) -> None:
     Depending on the use-case, the caller may need to clear any event listeners
     that are configured on the client.
     """
-    hello = client.admin.command(HelloCompat.LEGACY_CMD)
+    hello: dict = client.admin.command(HelloCompat.LEGACY_CMD)
     if "setName" not in hello:
         raise ConfigurationError("cluster is not a replica set")
 
@@ -612,7 +612,7 @@ def ensure_all_connected(client: MongoClient) -> None:
     def discover():
         i = 0
         while i < 100 and connected_host_list != target_host_list:
-            hello = client.admin.command(
+            hello: dict = client.admin.command(
                 HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY
             )
             connected_host_list.update([hello["me"]])