PYTHON-3493 Bulk Write InsertOne Should Be Parameter Of Collection Type (#1106)

Author: juliusgeo

doc/examples/type_hints.rst

@@ -113,6 +113,26 @@ These methods automatically add an "_id" field.
   >>> assert result is not None
   >>> assert result["year"] == 1993
   >>> # This will raise a type-checking error, despite being present, because it is added by PyMongo.
+  >>> assert result["_id"] # type:ignore[typeddict-item]
+
+This same typing scheme works for all of the insert methods (:meth:`~pymongo.collection.Collection.insert_one`,
+:meth:`~pymongo.collection.Collection.insert_many`, and :meth:`~pymongo.collection.Collection.bulk_write`).
+For `bulk_write` both :class:`~pymongo.operations.InsertOne` and :class:`~pymongo.operations.ReplaceOne` operators are generic.
+
+.. doctest::
+  :pyversion: >= 3.8
+
+  >>> from typing import TypedDict
+  >>> from pymongo import MongoClient
+  >>> from pymongo.operations import InsertOne
+  >>> from pymongo.collection import Collection
+  >>> client: MongoClient = MongoClient()
+  >>> collection: Collection[Movie] = client.test.test
+  >>> inserted = collection.bulk_write([InsertOne(Movie(name="Jurassic Park", year=1993))])
+  >>> result = collection.find_one({"name": "Jurassic Park"})
+  >>> assert result is not None
+  >>> assert result["year"] == 1993
+  >>> # This will raise a type-checking error, despite being present, because it is added by PyMongo.
   >>> assert result["_id"]  # type:ignore[typeddict-item]
 
 Modeling Document Types with TypedDict

mypy.ini

@@ -33,7 +33,7 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 
 [mypy-test.test_mypy]
-warn_unused_ignores = false
+warn_unused_ignores = True
 
 [mypy-winkerberos.*]
 ignore_missing_imports = True

pymongo/collection.py

@@ -77,7 +77,14 @@
 _FIND_AND_MODIFY_DOC_FIELDS = {"value": 1}
 
 
-_WriteOp = Union[InsertOne, DeleteOne, DeleteMany, ReplaceOne, UpdateOne, UpdateMany]
+_WriteOp = Union[
+    InsertOne[_DocumentType],
+    DeleteOne,
+    DeleteMany,
+    ReplaceOne[_DocumentType],
+    UpdateOne,
+    UpdateMany,
+]
 # Hint supports index name, "myIndex", or list of index pairs: [('x', 1), ('y', -1)]
 _IndexList = Sequence[Tuple[str, Union[int, str, Mapping[str, Any]]]]
 _IndexKeyHint = Union[str, _IndexList]
@@ -436,7 +443,7 @@ def with_options(
     @_csot.apply
     def bulk_write(
         self,
-        requests: Sequence[_WriteOp],
+        requests: Sequence[_WriteOp[_DocumentType]],
         ordered: bool = True,
         bypass_document_validation: bool = False,
         session: Optional["ClientSession"] = None,

pymongo/encryption.py

@@ -18,7 +18,7 @@
 import enum
 import socket
 import weakref
-from typing import Any, Mapping, Optional, Sequence
+from typing import Any, Generic, Mapping, Optional, Sequence
 
 try:
     from pymongocrypt.auto_encrypter import AutoEncrypter
@@ -55,6 +55,7 @@
 from pymongo.read_concern import ReadConcern
 from pymongo.results import BulkWriteResult, DeleteResult
 from pymongo.ssl_support import get_ssl_context
+from pymongo.typings import _DocumentType
 from pymongo.uri_parser import parse_host
 from pymongo.write_concern import WriteConcern
 
@@ -430,7 +431,7 @@ class QueryType(str, enum.Enum):
     """Used to encrypt a value for an equality query."""
 
 
-class ClientEncryption(object):
+class ClientEncryption(Generic[_DocumentType]):
     """Explicit client-side field level encryption."""
 
     def __init__(

pymongo/operations.py

@@ -13,21 +13,22 @@
 # limitations under the License.
 
 """Operation class definitions."""
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Generic, List, Mapping, Optional, Sequence, Tuple, Union
 
+from bson.raw_bson import RawBSONDocument
 from pymongo import helpers
 from pymongo.collation import validate_collation_or_none
 from pymongo.common import validate_boolean, validate_is_mapping, validate_list
 from pymongo.helpers import _gen_index_name, _index_document, _index_list
-from pymongo.typings import _CollationIn, _DocumentIn, _Pipeline
+from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
 
 
-class InsertOne(object):
+class InsertOne(Generic[_DocumentType]):
     """Represents an insert_one operation."""
 
     __slots__ = ("_doc",)
 
-    def __init__(self, document: _DocumentIn) -> None:
+    def __init__(self, document: Union[_DocumentType, RawBSONDocument]) -> None:
         """Create an InsertOne instance.
 
         For use with :meth:`~pymongo.collection.Collection.bulk_write`.
@@ -170,15 +171,15 @@ def __ne__(self, other: Any) -> bool:
         return not self == other
 
 
-class ReplaceOne(object):
+class ReplaceOne(Generic[_DocumentType]):
     """Represents a replace_one operation."""
 
     __slots__ = ("_filter", "_doc", "_upsert", "_collation", "_hint")
 
     def __init__(
         self,
         filter: Mapping[str, Any],
-        replacement: Mapping[str, Any],
+        replacement: Union[_DocumentType, RawBSONDocument],
         upsert: bool = False,
         collation: Optional[_CollationIn] = None,
         hint: Optional[_IndexKeyHint] = None,

pymongo/typings.py

@@ -37,3 +37,10 @@
 _Pipeline = Sequence[Mapping[str, Any]]
 _DocumentOut = _DocumentIn
 _DocumentType = TypeVar("_DocumentType", bound=Mapping[str, Any])
+
+
+def strip_optional(elem):
+    """This function is to allow us to cast all of the elements of an iterator from Optional[_T] to _T
+    while inside a list comprehension."""
+    assert elem is not None
+    return elem

test/__init__.py

@@ -1090,7 +1090,7 @@ def print_thread_stacks(pid: int) -> None:
 class IntegrationTest(PyMongoTestCase):
     """Base class for TestCases that need a connection to MongoDB to pass."""
 
-    client: MongoClient
+    client: MongoClient[dict]
     db: Database
     credentials: Dict[str, str]
 

test/mockupdb/test_cluster_time.py

@@ -60,7 +60,7 @@ def callback(client):
         self.cluster_time_conversation(callback, [{"ok": 1}] * 2)
 
     def test_bulk(self):
-        def callback(client):
+        def callback(client: MongoClient[dict]) -> None:
             client.db.collection.bulk_write(
                 [InsertOne({}), InsertOne({}), UpdateOne({}, {"$inc": {"x": 1}}), DeleteMany({})]
             )

test/mockupdb/test_op_msg.py

@@ -137,22 +137,22 @@
     # Legacy methods
     Operation(
         "bulk_write_insert",
-        lambda coll: coll.bulk_write([InsertOne({}), InsertOne({})]),
+        lambda coll: coll.bulk_write([InsertOne[dict]({}), InsertOne[dict]({})]),
         request=OpMsg({"insert": "coll"}, flags=0),
         reply={"ok": 1, "n": 2},
     ),
     Operation(
         "bulk_write_insert-w0",
         lambda coll: coll.with_options(write_concern=WriteConcern(w=0)).bulk_write(
-            [InsertOne({}), InsertOne({})]
+            [InsertOne[dict]({}), InsertOne[dict]({})]
         ),
         request=OpMsg({"insert": "coll"}, flags=0),
         reply={"ok": 1, "n": 2},
     ),
     Operation(
         "bulk_write_insert-w0-unordered",
         lambda coll: coll.with_options(write_concern=WriteConcern(w=0)).bulk_write(
-            [InsertOne({}), InsertOne({})], ordered=False
+            [InsertOne[dict]({}), InsertOne[dict]({})], ordered=False
         ),
         request=OpMsg({"insert": "coll"}, flags=OP_MSG_FLAGS["moreToCome"]),
         reply=None,

test/test_bulk.py

@@ -296,7 +296,7 @@ def test_upsert(self):
     def test_numerous_inserts(self):
         # Ensure we don't exceed server's maxWriteBatchSize size limit.
         n_docs = client_context.max_write_batch_size + 100
-        requests = [InsertOne({}) for _ in range(n_docs)]
+        requests = [InsertOne[dict]({}) for _ in range(n_docs)]
         result = self.coll.bulk_write(requests, ordered=False)
         self.assertEqual(n_docs, result.inserted_count)
         self.assertEqual(n_docs, self.coll.count_documents({}))
@@ -347,7 +347,7 @@ def test_bulk_write_no_results(self):
 
     def test_bulk_write_invalid_arguments(self):
         # The requests argument must be a list.
-        generator = (InsertOne({}) for _ in range(10))
+        generator = (InsertOne[dict]({}) for _ in range(10))
         with self.assertRaises(TypeError):
             self.coll.bulk_write(generator)  # type: ignore[arg-type]