PYTHON-3493 Bulk Write InsertOne Should Be Parameter Of Collection Type

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,
@@ -2044,7 +2051,7 @@ def create_index(
             cmd_options["maxTimeMS"] = kwargs.pop("maxTimeMS")
         if comment is not None:
             cmd_options["comment"] = comment
-        index = IndexModel(keys, **kwargs)
+        index: IndexModel[_DocumentType] = IndexModel(keys, **kwargs)
         return self.__create_indexes([index], session, **cmd_options)[0]
 
     def drop_indexes(

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(object, Generic[_DocumentType]):
     """Explicit client-side field level encryption."""
 
     def __init__(

pymongo/operations.py

@@ -13,21 +13,21 @@
 # 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 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 RawBSONDocument, _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`.
@@ -58,7 +58,7 @@ def __ne__(self, other: Any) -> bool:
 _IndexKeyHint = Union[str, _IndexList]
 
 
-class DeleteOne(object):
+class DeleteOne(object, Generic[_DocumentType]):
     """Represents a delete_one operation."""
 
     __slots__ = ("_filter", "_collation", "_hint")
@@ -170,15 +170,15 @@ def __ne__(self, other: Any) -> bool:
         return not self == other
 
 
-class ReplaceOne(object):
+class ReplaceOne(object, 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,
@@ -308,15 +308,15 @@ def __repr__(self):
         )
 
 
-class UpdateOne(_UpdateOp):
+class UpdateOne(_UpdateOp, Generic[_DocumentType]):
     """Represents an update_one operation."""
 
     __slots__ = ()
 
     def __init__(
         self,
         filter: Mapping[str, Any],
-        update: Union[Mapping[str, Any], _Pipeline],
+        update: Union[_DocumentType, RawBSONDocument, _Pipeline],
         upsert: bool = False,
         collation: Optional[_CollationIn] = None,
         array_filters: Optional[List[Mapping[str, Any]]] = None,
@@ -366,15 +366,15 @@ def _add_to_bulk(self, bulkobj):
         )
 
 
-class UpdateMany(_UpdateOp):
+class UpdateMany(_UpdateOp, Generic[_DocumentType]):
     """Represents an update_many operation."""
 
     __slots__ = ()
 
     def __init__(
         self,
         filter: Mapping[str, Any],
-        update: Union[Mapping[str, Any], _Pipeline],
+        update: Union[_DocumentType, RawBSONDocument, _Pipeline],
         upsert: bool = False,
         collation: Optional[_CollationIn] = None,
         array_filters: Optional[List[Mapping[str, Any]]] = None,
@@ -424,7 +424,7 @@ def _add_to_bulk(self, bulkobj):
         )
 
 
-class IndexModel(object):
+class IndexModel(object, Generic[_DocumentType]):
     """Represents an index to create."""
 
     __slots__ = ("__document",)

test/test_mypy.py

@@ -22,7 +22,14 @@
 try:
     from typing_extensions import TypedDict
 
-    class Movie(TypedDict):  # type: ignore[misc]
+    from bson import ObjectId
+
+    class Movie(TypedDict):
+        name: str
+        year: int
+
+    class MovieWithId(TypedDict):
+        _id: ObjectId
         name: str
         year: int
 
@@ -43,7 +50,7 @@ class Movie(TypedDict):  # type: ignore[misc]
 from bson.son import SON
 from pymongo import ASCENDING, MongoClient
 from pymongo.collection import Collection
-from pymongo.operations import InsertOne
+from pymongo.operations import InsertOne, ReplaceOne
 from pymongo.read_preferences import ReadPreference
 
 TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "mypy_fails")
@@ -110,8 +117,9 @@ def to_list(iterable: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]:
 
     def test_bulk_write(self) -> None:
         self.coll.insert_one({})
-        requests = [InsertOne({})]
-        result = self.coll.bulk_write(requests)
+        coll: Collection[Movie] = self.coll
+        requests: List[InsertOne[Movie]] = [InsertOne(Movie(name="American Graffiti", year=1973))]
+        result = coll.bulk_write(requests)
         self.assertTrue(result.acknowledged)
 
     def test_command(self) -> None:
@@ -324,6 +332,40 @@ def test_typeddict_document_type_insertion(self) -> None:
         )
         coll.insert_many([bad_movie])
 
+    @only_type_check
+    def test_bulk_write_document_type_insertion(self):
+        client: MongoClient[MovieWithId] = MongoClient()
+        coll: Collection[MovieWithId] = client.test.test
+        coll.bulk_write(
+            [InsertOne(Movie({"name": "THX-1138", "year": 1971}))]  # type:ignore[arg-type]
+        )
+        mov_dict = {"_id": ObjectId(), "name": "THX-1138", "year": 1971}
+        coll.bulk_write(
+            [InsertOne(mov_dict)]  # type:ignore[arg-type]
+        )
+        coll.bulk_write(
+            [
+                InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971})
+            ]  # No error because it is in-line.
+        )
+
+    @only_type_check
+    def test_bulk_write_document_type_replacement(self):
+        client: MongoClient[MovieWithId] = MongoClient()
+        coll: Collection[MovieWithId] = client.test.test
+        coll.bulk_write(
+            [ReplaceOne({}, Movie({"name": "THX-1138", "year": 1971}))]  # type:ignore[arg-type]
+        )
+        mov_dict = {"_id": ObjectId(), "name": "THX-1138", "year": 1971}
+        coll.bulk_write(
+            [ReplaceOne({}, mov_dict)]  # type:ignore[arg-type]
+        )
+        coll.bulk_write(
+            [
+                ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971})
+            ]  # No error because it is in-line.
+        )
+
     @only_type_check
     def test_raw_bson_document_type(self) -> None:
         client = MongoClient(document_class=RawBSONDocument)