Skip to content

Commit 3ef565f

Browse files
blink1073blink1073
authored
PYTHON-4796 Update type checkers and handle with_options typing (#1880)
1 parent 1e395de commit 3ef565f

File tree

20 files changed

+132
-33
lines changed

20 files changed

+132
-33
lines changed

‎bson/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1324,7 +1324,7 @@ def decode_iter(
13241324
elements = data[position : position + obj_size]
13251325
position += obj_size
13261326

1327-
yield _bson_to_dict(elements, opts) # type:ignore[misc, type-var]
1327+
yield _bson_to_dict(elements, opts) # type:ignore[misc]
13281328

13291329

13301330
@overload
@@ -1370,7 +1370,7 @@ def decode_file_iter(
13701370
raise InvalidBSON("cut off in middle of objsize")
13711371
obj_size = _UNPACK_INT_FROM(size_data, 0)[0] - 4
13721372
elements = size_data + file_obj.read(max(0, obj_size))
1373-
yield _bson_to_dict(elements, opts) # type:ignore[type-var, arg-type, misc]
1373+
yield _bson_to_dict(elements, opts) # type:ignore[arg-type, misc]
13741374

13751375

13761376
def is_valid(bson: bytes) -> bool:

‎bson/decimal128.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ def __init__(self, value: _VALUE_OPTIONS) -> None:
223223
"from list or tuple. Must have exactly 2 "
224224
"elements."
225225
)
226-
self.__high, self.__low = value # type: ignore
226+
self.__high, self.__low = value
227227
else:
228228
raise TypeError(f"Cannot convert {value!r} to Decimal128")
229229

‎bson/json_util.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ def __new__(
324324
"JSONOptions.datetime_representation must be one of LEGACY, "
325325
"NUMBERLONG, or ISO8601 from DatetimeRepresentation."
326326
)
327-
self = cast(JSONOptions, super().__new__(cls, *args, **kwargs)) # type:ignore[arg-type]
327+
self = cast(JSONOptions, super().__new__(cls, *args, **kwargs))
328328
if json_mode not in (JSONMode.LEGACY, JSONMode.RELAXED, JSONMode.CANONICAL):
329329
raise ValueError(
330330
"JSONOptions.json_mode must be one of LEGACY, RELAXED, "

‎bson/son.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ def __init__(
6868
self.update(kwargs)
6969

7070
def __new__(cls: Type[SON[_Key, _Value]], *args: Any, **kwargs: Any) -> SON[_Key, _Value]:
71-
instance = super().__new__(cls, *args, **kwargs) # type: ignore[type-var]
71+
instance = super().__new__(cls, *args, **kwargs)
7272
instance.__keys = []
7373
return instance
7474

‎hatch.toml

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,9 @@ features = ["docs","test"]
1313
test = "sphinx-build -E -b doctest doc ./doc/_build/doctest"
1414

1515
[envs.typing]
16-
features = ["encryption", "ocsp", "zstd", "aws"]
17-
dependencies = ["mypy==1.2.0","pyright==1.1.290", "certifi", "typing_extensions"]
16+
pre-install-commands = [
17+
"pip install -q -r requirements/typing.txt",
18+
]
1819
[envs.typing.scripts]
1920
check-mypy = [
2021
"mypy --install-types --non-interactive bson gridfs tools pymongo",

‎pymongo/_csot.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,14 +75,13 @@ def __init__(self, timeout: Optional[float]):
7575
self._timeout = timeout
7676
self._tokens: Optional[tuple[Token[Optional[float]], Token[float], Token[float]]] = None
7777

78-
def __enter__(self) -> _TimeoutContext:
78+
def __enter__(self) -> None:
7979
timeout_token = TIMEOUT.set(self._timeout)
8080
prev_deadline = DEADLINE.get()
8181
next_deadline = time.monotonic() + self._timeout if self._timeout else float("inf")
8282
deadline_token = DEADLINE.set(min(prev_deadline, next_deadline))
8383
rtt_token = RTT.set(0.0)
8484
self._tokens = (timeout_token, deadline_token, rtt_token)
85-
return self
8685

8786
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
8887
if self._tokens:

‎pymongo/asynchronous/collection.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
TypeVar,
3636
Union,
3737
cast,
38+
overload,
3839
)
3940

4041
from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
@@ -332,13 +333,33 @@ def database(self) -> AsyncDatabase[_DocumentType]:
332333
"""
333334
return self._database
334335

336+
@overload
337+
def with_options(
338+
self,
339+
codec_options: None = None,
340+
read_preference: Optional[_ServerMode] = ...,
341+
write_concern: Optional[WriteConcern] = ...,
342+
read_concern: Optional[ReadConcern] = ...,
343+
) -> AsyncCollection[_DocumentType]:
344+
...
345+
346+
@overload
347+
def with_options(
348+
self,
349+
codec_options: bson.CodecOptions[_DocumentTypeArg],
350+
read_preference: Optional[_ServerMode] = ...,
351+
write_concern: Optional[WriteConcern] = ...,
352+
read_concern: Optional[ReadConcern] = ...,
353+
) -> AsyncCollection[_DocumentTypeArg]:
354+
...
355+
335356
def with_options(
336357
self,
337358
codec_options: Optional[bson.CodecOptions[_DocumentTypeArg]] = None,
338359
read_preference: Optional[_ServerMode] = None,
339360
write_concern: Optional[WriteConcern] = None,
340361
read_concern: Optional[ReadConcern] = None,
341-
) -> AsyncCollection[_DocumentType]:
362+
) -> AsyncCollection[_DocumentType] | AsyncCollection[_DocumentTypeArg]:
342363
"""Get a clone of this collection changing the specified settings.
343364
344365
>>> coll1.read_preference

‎pymongo/asynchronous/database.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,13 +146,33 @@ def name(self) -> str:
146146
"""The name of this :class:`AsyncDatabase`."""
147147
return self._name
148148

149+
@overload
150+
def with_options(
151+
self,
152+
codec_options: None = None,
153+
read_preference: Optional[_ServerMode] = ...,
154+
write_concern: Optional[WriteConcern] = ...,
155+
read_concern: Optional[ReadConcern] = ...,
156+
) -> AsyncDatabase[_DocumentType]:
157+
...
158+
159+
@overload
160+
def with_options(
161+
self,
162+
codec_options: bson.CodecOptions[_DocumentTypeArg],
163+
read_preference: Optional[_ServerMode] = ...,
164+
write_concern: Optional[WriteConcern] = ...,
165+
read_concern: Optional[ReadConcern] = ...,
166+
) -> AsyncDatabase[_DocumentTypeArg]:
167+
...
168+
149169
def with_options(
150170
self,
151171
codec_options: Optional[CodecOptions[_DocumentTypeArg]] = None,
152172
read_preference: Optional[_ServerMode] = None,
153173
write_concern: Optional[WriteConcern] = None,
154174
read_concern: Optional[ReadConcern] = None,
155-
) -> AsyncDatabase[_DocumentType]:
175+
) -> AsyncDatabase[_DocumentType] | AsyncDatabase[_DocumentTypeArg]:
156176
"""Get a clone of this database changing the specified settings.
157177
158178
>>> db1.read_preference

‎pymongo/asynchronous/pool.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -913,7 +913,7 @@ async def _configured_socket(
913913
and not options.tls_allow_invalid_hostnames
914914
):
915915
try:
916-
ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)
916+
ssl.match_hostname(ssl_sock.getpeercert(), hostname=host) # type:ignore[attr-defined]
917917
except _CertificateError:
918918
ssl_sock.close()
919919
raise

‎pymongo/common.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -850,7 +850,7 @@ def get_normed_key(x: str) -> str:
850850
return x
851851

852852
def get_setter_key(x: str) -> str:
853-
return options.cased_key(x) # type: ignore[attr-defined]
853+
return options.cased_key(x)
854854

855855
else:
856856
validated_options = {}

‎pymongo/compression_support.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626

2727
def _have_snappy() -> bool:
2828
try:
29-
import snappy # type:ignore[import] # noqa: F401
29+
import snappy # type:ignore[import-not-found] # noqa: F401
3030

3131
return True
3232
except ImportError:

‎pymongo/encryption_options.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from typing import TYPE_CHECKING, Any, Mapping, Optional
2222

2323
try:
24-
import pymongocrypt # type:ignore[import] # noqa: F401
24+
import pymongocrypt # type:ignore[import-untyped] # noqa: F401
2525

2626
# Check for pymongocrypt>=1.10.
2727
from pymongocrypt import synchronous as _ # noqa: F401

‎pymongo/synchronous/collection.py

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
TypeVar,
3535
Union,
3636
cast,
37+
overload,
3738
)
3839

3940
from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
@@ -333,13 +334,33 @@ def database(self) -> Database[_DocumentType]:
333334
"""
334335
return self._database
335336

337+
@overload
338+
def with_options(
339+
self,
340+
codec_options: None = None,
341+
read_preference: Optional[_ServerMode] = ...,
342+
write_concern: Optional[WriteConcern] = ...,
343+
read_concern: Optional[ReadConcern] = ...,
344+
) -> Collection[_DocumentType]:
345+
...
346+
347+
@overload
348+
def with_options(
349+
self,
350+
codec_options: bson.CodecOptions[_DocumentTypeArg],
351+
read_preference: Optional[_ServerMode] = ...,
352+
write_concern: Optional[WriteConcern] = ...,
353+
read_concern: Optional[ReadConcern] = ...,
354+
) -> Collection[_DocumentTypeArg]:
355+
...
356+
336357
def with_options(
337358
self,
338359
codec_options: Optional[bson.CodecOptions[_DocumentTypeArg]] = None,
339360
read_preference: Optional[_ServerMode] = None,
340361
write_concern: Optional[WriteConcern] = None,
341362
read_concern: Optional[ReadConcern] = None,
342-
) -> Collection[_DocumentType]:
363+
) -> Collection[_DocumentType] | Collection[_DocumentTypeArg]:
343364
"""Get a clone of this collection changing the specified settings.
344365
345366
>>> coll1.read_preference

‎pymongo/synchronous/database.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,13 +146,33 @@ def name(self) -> str:
146146
"""The name of this :class:`Database`."""
147147
return self._name
148148

149+
@overload
150+
def with_options(
151+
self,
152+
codec_options: None = None,
153+
read_preference: Optional[_ServerMode] = ...,
154+
write_concern: Optional[WriteConcern] = ...,
155+
read_concern: Optional[ReadConcern] = ...,
156+
) -> Database[_DocumentType]:
157+
...
158+
159+
@overload
160+
def with_options(
161+
self,
162+
codec_options: bson.CodecOptions[_DocumentTypeArg],
163+
read_preference: Optional[_ServerMode] = ...,
164+
write_concern: Optional[WriteConcern] = ...,
165+
read_concern: Optional[ReadConcern] = ...,
166+
) -> Database[_DocumentTypeArg]:
167+
...
168+
149169
def with_options(
150170
self,
151171
codec_options: Optional[CodecOptions[_DocumentTypeArg]] = None,
152172
read_preference: Optional[_ServerMode] = None,
153173
write_concern: Optional[WriteConcern] = None,
154174
read_concern: Optional[ReadConcern] = None,
155-
) -> Database[_DocumentType]:
175+
) -> Database[_DocumentType] | Database[_DocumentTypeArg]:
156176
"""Get a clone of this database changing the specified settings.
157177
158178
>>> db1.read_preference

‎pymongo/synchronous/pool.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,7 @@ def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.
909909
and not options.tls_allow_invalid_hostnames
910910
):
911911
try:
912-
ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)
912+
ssl.match_hostname(ssl_sock.getpeercert(), hostname=host) # type:ignore[attr-defined]
913913
except _CertificateError:
914914
ssl_sock.close()
915915
raise

‎requirements/typing.txt

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
mypy==1.11.2
2+
pyright==1.1.382.post1
3+
typing_extensions
4+
-r ./encryption.txt
5+
-r ./ocsp.txt
6+
-r ./zstd.txt
7+
-r ./aws.txt

‎test/asynchronous/test_database.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -711,7 +711,7 @@ def test_with_options(self):
711711
"write_concern": WriteConcern(w=1),
712712
"read_concern": ReadConcern(level="local"),
713713
}
714-
db2 = db1.with_options(**newopts) # type: ignore[arg-type]
714+
db2 = db1.with_options(**newopts) # type: ignore[arg-type, call-overload]
715715
for opt in newopts:
716716
self.assertEqual(getattr(db2, opt), newopts.get(opt, getattr(db1, opt)))
717717

‎test/test_database.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -702,7 +702,7 @@ def test_with_options(self):
702702
"write_concern": WriteConcern(w=1),
703703
"read_concern": ReadConcern(level="local"),
704704
}
705-
db2 = db1.with_options(**newopts) # type: ignore[arg-type]
705+
db2 = db1.with_options(**newopts) # type: ignore[arg-type, call-overload]
706706
for opt in newopts:
707707
self.assertEqual(getattr(db2, opt), newopts.get(opt, getattr(db1, opt)))
708708

‎test/test_typing.py

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
cast,
3535
)
3636

37-
try:
37+
if TYPE_CHECKING:
3838
from typing_extensions import NotRequired, TypedDict
3939

4040
from bson import ObjectId
@@ -49,16 +49,13 @@ class MovieWithId(TypedDict):
4949
year: int
5050

5151
class ImplicitMovie(TypedDict):
52-
_id: NotRequired[ObjectId] # pyright: ignore[reportGeneralTypeIssues]
52+
_id: NotRequired[ObjectId]
5353
name: str
5454
year: int
55-
56-
except ImportError:
57-
Movie = dict # type:ignore[misc,assignment]
58-
ImplicitMovie = dict # type: ignore[assignment,misc]
59-
MovieWithId = dict # type: ignore[assignment,misc]
60-
TypedDict = None
61-
NotRequired = None # type: ignore[assignment]
55+
else:
56+
Movie = dict
57+
ImplicitMovie = dict
58+
NotRequired = None
6259

6360

6461
try:
@@ -234,6 +231,19 @@ def execute_transaction(session):
234231
execute_transaction, read_preference=ReadPreference.PRIMARY
235232
)
236233

234+
def test_with_options(self) -> None:
235+
coll: Collection[Dict[str, Any]] = self.coll
236+
coll.drop()
237+
doc = {"name": "foo", "year": 1982, "other": 1}
238+
coll.insert_one(doc)
239+
240+
coll2 = coll.with_options(codec_options=CodecOptions(document_class=Movie))
241+
retrieved = coll2.find_one()
242+
assert retrieved is not None
243+
assert retrieved["name"] == "foo"
244+
# We expect a type error here.
245+
assert retrieved["other"] == 1 # type:ignore[typeddict-item]
246+
237247

238248
class TestDecode(unittest.TestCase):
239249
def test_bson_decode(self) -> None:
@@ -426,7 +436,7 @@ def test_bulk_write_document_type_insertion(self):
426436
)
427437
coll.bulk_write(
428438
[
429-
InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971})
439+
InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971}) # pyright: ignore
430440
] # No error because it is in-line.
431441
)
432442

@@ -443,7 +453,7 @@ def test_bulk_write_document_type_replacement(self):
443453
)
444454
coll.bulk_write(
445455
[
446-
ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971})
456+
ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971}) # pyright: ignore
447457
] # No error because it is in-line.
448458
)
449459

@@ -566,7 +576,7 @@ def test_explicit_document_type(self) -> None:
566576
def test_typeddict_document_type(self) -> None:
567577
options: CodecOptions[Movie] = CodecOptions()
568578
# Suppress: Cannot instantiate type "Type[Movie]".
569-
obj = options.document_class(name="a", year=1) # type: ignore[misc]
579+
obj = options.document_class(name="a", year=1)
570580
assert obj["year"] == 1
571581
assert obj["name"] == "a"
572582

‎tools/synchro.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
from os import listdir
2424
from pathlib import Path
2525

26-
from unasync import Rule, unasync_files # type: ignore[import]
26+
from unasync import Rule, unasync_files # type: ignore[import-not-found]
2727

2828
replacements = {
2929
"AsyncCollection": "Collection",

0 commit comments

Comments
 (0)