Skip to content

Commit 9485e42

Browse files
committed
fix typing -- i was being so silly earlier
1 parent 19bcce4 commit 9485e42

File tree

5 files changed

+26
-22
lines changed

5 files changed

+26
-22
lines changed

pymongo/asynchronous/bulk.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,9 @@ def add_delete(
243243
def gen_ordered(
244244
self,
245245
requests: Iterable[Any],
246-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
246+
process: Union[
247+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
248+
],
247249
) -> Iterator[_Run]:
248250
"""Generate batches of operations, batched by type of
249251
operation, in the order **provided**.
@@ -266,7 +268,9 @@ def gen_ordered(
266268
def gen_unordered(
267269
self,
268270
requests: Iterable[Any],
269-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
271+
process: Union[
272+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
273+
],
270274
) -> Iterator[_Run]:
271275
"""Generate batches of operations, batched by type of
272276
operation, in arbitrary order.
@@ -787,7 +791,9 @@ async def validate_batch(self, conn: AsyncConnection, write_concern: WriteConcer
787791
async def execute(
788792
self,
789793
generator: Iterable[Any],
790-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
794+
process: Union[
795+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
796+
],
791797
write_concern: WriteConcern,
792798
session: Optional[AsyncClientSession],
793799
operation: str,

pymongo/asynchronous/collection.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -785,7 +785,7 @@ async def bulk_write(
785785

786786
write_concern = self._write_concern_for(session)
787787

788-
def process_for_bulk(request: Union[_DocumentType, RawBSONDocument, _WriteOp]) -> bool:
788+
def process_for_bulk(request: _WriteOp) -> bool:
789789
try:
790790
return request._add_to_bulk(blk)
791791
except AttributeError:
@@ -964,7 +964,7 @@ async def insert_many(
964964
raise TypeError("documents must be a non-empty list")
965965
inserted_ids: list[ObjectId] = []
966966

967-
def process_for_bulk(document: Union[_DocumentType, RawBSONDocument, _WriteOp]) -> bool:
967+
def process_for_bulk(document: Union[_DocumentType, RawBSONDocument]) -> bool:
968968
"""A generator that validates documents and handles _ids."""
969969
common.validate_is_document_type("document", document)
970970
if not isinstance(document, RawBSONDocument):

pymongo/synchronous/bulk.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,9 @@ def add_delete(
243243
def gen_ordered(
244244
self,
245245
requests: Iterable[Any],
246-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
246+
process: Union[
247+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
248+
],
247249
) -> Iterator[_Run]:
248250
"""Generate batches of operations, batched by type of
249251
operation, in the order **provided**.
@@ -266,7 +268,9 @@ def gen_ordered(
266268
def gen_unordered(
267269
self,
268270
requests: Iterable[Any],
269-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
271+
process: Union[
272+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
273+
],
270274
) -> Iterator[_Run]:
271275
"""Generate batches of operations, batched by type of
272276
operation, in arbitrary order.
@@ -787,7 +791,9 @@ def validate_batch(self, conn: Connection, write_concern: WriteConcern) -> None:
787791
def execute(
788792
self,
789793
generator: Iterable[Any],
790-
process: Callable[[Union[_DocumentType, RawBSONDocument, _WriteOp]], bool],
794+
process: Union[
795+
Callable[[_WriteOp], bool], Callable[[Union[_DocumentType, RawBSONDocument]], bool]
796+
],
791797
write_concern: WriteConcern,
792798
session: Optional[ClientSession],
793799
operation: str,

pymongo/synchronous/collection.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -784,7 +784,7 @@ def bulk_write(
784784

785785
write_concern = self._write_concern_for(session)
786786

787-
def process_for_bulk(request: Union[_DocumentType, RawBSONDocument, _WriteOp]) -> bool:
787+
def process_for_bulk(request: _WriteOp) -> bool:
788788
try:
789789
return request._add_to_bulk(blk)
790790
except AttributeError:
@@ -963,7 +963,7 @@ def insert_many(
963963
raise TypeError("documents must be a non-empty list")
964964
inserted_ids: list[ObjectId] = []
965965

966-
def process_for_bulk(document: Union[_DocumentType, RawBSONDocument, _WriteOp]) -> bool:
966+
def process_for_bulk(document: Union[_DocumentType, RawBSONDocument]) -> bool:
967967
"""A generator that validates documents and handles _ids."""
968968
common.validate_is_document_type("document", document)
969969
if not isinstance(document, RawBSONDocument):

test/test_typing.py

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -428,13 +428,9 @@ def test_typeddict_document_type_insertion(self) -> None:
428428
def test_bulk_write_document_type_insertion(self):
429429
client: MongoClient[MovieWithId] = MongoClient()
430430
coll: Collection[MovieWithId] = client.test.test
431-
coll.bulk_write(
432-
[InsertOne(Movie({"name": "THX-1138", "year": 1971}))] # type:ignore[arg-type]
433-
)
431+
coll.bulk_write([InsertOne(Movie({"name": "THX-1138", "year": 1971}))])
434432
mov_dict = {"_id": ObjectId(), "name": "THX-1138", "year": 1971}
435-
coll.bulk_write(
436-
[InsertOne(mov_dict)] # type:ignore[arg-type]
437-
)
433+
coll.bulk_write([InsertOne(mov_dict)])
438434
coll.bulk_write(
439435
[
440436
InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971}) # pyright: ignore
@@ -445,13 +441,9 @@ def test_bulk_write_document_type_insertion(self):
445441
def test_bulk_write_document_type_replacement(self):
446442
client: MongoClient[MovieWithId] = MongoClient()
447443
coll: Collection[MovieWithId] = client.test.test
448-
coll.bulk_write(
449-
[ReplaceOne({}, Movie({"name": "THX-1138", "year": 1971}))] # type:ignore[arg-type]
450-
)
444+
coll.bulk_write([ReplaceOne({}, Movie({"name": "THX-1138", "year": 1971}))])
451445
mov_dict = {"_id": ObjectId(), "name": "THX-1138", "year": 1971}
452-
coll.bulk_write(
453-
[ReplaceOne({}, mov_dict)] # type:ignore[arg-type]
454-
)
446+
coll.bulk_write([ReplaceOne({}, mov_dict)])
455447
coll.bulk_write(
456448
[
457449
ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971}) # pyright: ignore

0 commit comments

Comments
 (0)