PYTHON-4796 Update type checkers and handle with_options typing (#1880)

bson/__init__.py

@@ -1324,7 +1324,7 @@ def decode_iter(
         elements = data[position : position + obj_size]
         position += obj_size
 
-        yield _bson_to_dict(elements, opts)  # type:ignore[misc, type-var]
+        yield _bson_to_dict(elements, opts)  # type:ignore[misc]
 
 
 @overload
@@ -1370,7 +1370,7 @@ def decode_file_iter(
             raise InvalidBSON("cut off in middle of objsize")
         obj_size = _UNPACK_INT_FROM(size_data, 0)[0] - 4
         elements = size_data + file_obj.read(max(0, obj_size))
-        yield _bson_to_dict(elements, opts)  # type:ignore[type-var, arg-type, misc]
+        yield _bson_to_dict(elements, opts)  # type:ignore[arg-type, misc]
 
 
 def is_valid(bson: bytes) -> bool:

bson/decimal128.py

@@ -223,7 +223,7 @@ class Decimal128:
                     "from list or tuple. Must have exactly 2 "
                     "elements."
                 )
-            self.__high, self.__low = value  # type: ignore
+            self.__high, self.__low = value
         else:
             raise TypeError(f"Cannot convert {value!r} to Decimal128")
 

bson/json_util.py

@@ -324,7 +324,7 @@ class JSONOptions(_BASE_CLASS):
                 "JSONOptions.datetime_representation must be one of LEGACY, "
                 "NUMBERLONG, or ISO8601 from DatetimeRepresentation."
             )
-        self = cast(JSONOptions, super().__new__(cls, *args, **kwargs))  # type:ignore[arg-type]
+        self = cast(JSONOptions, super().__new__(cls, *args, **kwargs))
         if json_mode not in (JSONMode.LEGACY, JSONMode.RELAXED, JSONMode.CANONICAL):
             raise ValueError(
                 "JSONOptions.json_mode must be one of LEGACY, RELAXED, "

bson/son.py

@@ -68,7 +68,7 @@ class SON(Dict[_Key, _Value]):
         self.update(kwargs)
 
     def __new__(cls: Type[SON[_Key, _Value]], *args: Any, **kwargs: Any) -> SON[_Key, _Value]:
-        instance = super().__new__(cls, *args, **kwargs)  # type: ignore[type-var]
+        instance = super().__new__(cls, *args, **kwargs)
         instance.__keys = []
         return instance
 

hatch.toml

@@ -13,8 +13,9 @@ features = ["docs","test"]
 test = "sphinx-build -E -b doctest doc ./doc/_build/doctest"
 
 [envs.typing]
-features = ["encryption", "ocsp", "zstd", "aws"]
-dependencies = ["mypy==1.2.0","pyright==1.1.290", "certifi", "typing_extensions"]
+pre-install-commands = [
+  "pip install -q -r requirements/typing.txt",
+]
 [envs.typing.scripts]
 check-mypy = [
     "mypy --install-types --non-interactive bson gridfs tools pymongo",

pymongo/_csot.py

@@ -75,14 +75,13 @@ class _TimeoutContext(AbstractContextManager):
         self._timeout = timeout
         self._tokens: Optional[tuple[Token[Optional[float]], Token[float], Token[float]]] = None
 
-    def __enter__(self) -> _TimeoutContext:
+    def __enter__(self) -> None:
         timeout_token = TIMEOUT.set(self._timeout)
         prev_deadline = DEADLINE.get()
         next_deadline = time.monotonic() + self._timeout if self._timeout else float("inf")
         deadline_token = DEADLINE.set(min(prev_deadline, next_deadline))
         rtt_token = RTT.set(0.0)
         self._tokens = (timeout_token, deadline_token, rtt_token)
-        return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if self._tokens:

pymongo/asynchronous/collection.py

@@ -35,6 +35,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
@@ -332,13 +333,33 @@ class AsyncCollection(common.BaseObject, Generic[_DocumentType]):
         """
         return self._database
 
+    @overload
+    def with_options(
+        self,
+        codec_options: None = None,
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> AsyncCollection[_DocumentType]:
+        ...
+
+    @overload
+    def with_options(
+        self,
+        codec_options: bson.CodecOptions[_DocumentTypeArg],
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> AsyncCollection[_DocumentTypeArg]:
+        ...
+
     def with_options(
         self,
         codec_options: Optional[bson.CodecOptions[_DocumentTypeArg]] = None,
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> AsyncCollection[_DocumentType]:
+    ) -> AsyncCollection[_DocumentType] | AsyncCollection[_DocumentTypeArg]:
         """Get a clone of this collection changing the specified settings.
 
           >>> coll1.read_preference

pymongo/asynchronous/database.py

@@ -146,13 +146,33 @@ class AsyncDatabase(common.BaseObject, Generic[_DocumentType]):
         """The name of this :class:`AsyncDatabase`."""
         return self._name
 
+    @overload
+    def with_options(
+        self,
+        codec_options: None = None,
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> AsyncDatabase[_DocumentType]:
+        ...
+
+    @overload
+    def with_options(
+        self,
+        codec_options: bson.CodecOptions[_DocumentTypeArg],
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> AsyncDatabase[_DocumentTypeArg]:
+        ...
+
     def with_options(
         self,
         codec_options: Optional[CodecOptions[_DocumentTypeArg]] = None,
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> AsyncDatabase[_DocumentType]:
+    ) -> AsyncDatabase[_DocumentType] | AsyncDatabase[_DocumentTypeArg]:
         """Get a clone of this database changing the specified settings.
 
           >>> db1.read_preference

pymongo/asynchronous/pool.py

@@ -913,7 +913,7 @@ async def _configured_socket(
         and not options.tls_allow_invalid_hostnames
     ):
         try:
-            ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)
+            ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)  # type:ignore[attr-defined]
         except _CertificateError:
             ssl_sock.close()
             raise

pymongo/common.py

@@ -850,7 +850,7 @@ def get_validated_options(
             return x
 
         def get_setter_key(x: str) -> str:
-            return options.cased_key(x)  # type: ignore[attr-defined]
+            return options.cased_key(x)
 
     else:
         validated_options = {}

pymongo/compression_support.py

@@ -26,7 +26,7 @@ _NO_COMPRESSION.update(_SENSITIVE_COMMANDS)
 
 def _have_snappy() -> bool:
     try:
-        import snappy  # type:ignore[import]  # noqa: F401
+        import snappy  # type:ignore[import-not-found]  # noqa: F401
 
         return True
     except ImportError:

pymongo/encryption_options.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Mapping, Optional
 
 try:
-    import pymongocrypt  # type:ignore[import] # noqa: F401
+    import pymongocrypt  # type:ignore[import-untyped] # noqa: F401
 
     # Check for pymongocrypt>=1.10.
     from pymongocrypt import synchronous as _  # noqa: F401

pymongo/synchronous/collection.py

@@ -34,6 +34,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
@@ -333,13 +334,33 @@ class Collection(common.BaseObject, Generic[_DocumentType]):
         """
         return self._database
 
+    @overload
+    def with_options(
+        self,
+        codec_options: None = None,
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> Collection[_DocumentType]:
+        ...
+
+    @overload
+    def with_options(
+        self,
+        codec_options: bson.CodecOptions[_DocumentTypeArg],
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> Collection[_DocumentTypeArg]:
+        ...
+
     def with_options(
         self,
         codec_options: Optional[bson.CodecOptions[_DocumentTypeArg]] = None,
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> Collection[_DocumentType]:
+    ) -> Collection[_DocumentType] | Collection[_DocumentTypeArg]:
         """Get a clone of this collection changing the specified settings.
 
           >>> coll1.read_preference

pymongo/synchronous/database.py

@@ -146,13 +146,33 @@ class Database(common.BaseObject, Generic[_DocumentType]):
         """The name of this :class:`Database`."""
         return self._name
 
+    @overload
+    def with_options(
+        self,
+        codec_options: None = None,
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> Database[_DocumentType]:
+        ...
+
+    @overload
+    def with_options(
+        self,
+        codec_options: bson.CodecOptions[_DocumentTypeArg],
+        read_preference: Optional[_ServerMode] = ...,
+        write_concern: Optional[WriteConcern] = ...,
+        read_concern: Optional[ReadConcern] = ...,
+    ) -> Database[_DocumentTypeArg]:
+        ...
+
     def with_options(
         self,
         codec_options: Optional[CodecOptions[_DocumentTypeArg]] = None,
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> Database[_DocumentType]:
+    ) -> Database[_DocumentType] | Database[_DocumentTypeArg]:
         """Get a clone of this database changing the specified settings.
 
           >>> db1.read_preference

pymongo/synchronous/pool.py

@@ -909,7 +909,7 @@ def _configured_socket(address: _Address, options: PoolOptions) -> Union[socket.
         and not options.tls_allow_invalid_hostnames
     ):
         try:
-            ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)
+            ssl.match_hostname(ssl_sock.getpeercert(), hostname=host)  # type:ignore[attr-defined]
         except _CertificateError:
             ssl_sock.close()
             raise

requirements/typing.txt

@@ -0,0 +1,7 @@
+mypy==1.11.2
+pyright==1.1.382.post1
+typing_extensions
+-r ./encryption.txt
+-r ./ocsp.txt
+-r ./zstd.txt
+-r ./aws.txt

test/asynchronous/test_database.py

@@ -711,7 +711,7 @@ class TestDatabase(AsyncIntegrationTest):
             "write_concern": WriteConcern(w=1),
             "read_concern": ReadConcern(level="local"),
         }
-        db2 = db1.with_options(**newopts)  # type: ignore[arg-type]
+        db2 = db1.with_options(**newopts)  # type: ignore[arg-type, call-overload]
         for opt in newopts:
             self.assertEqual(getattr(db2, opt), newopts.get(opt, getattr(db1, opt)))
 

test/test_database.py

@@ -702,7 +702,7 @@ class TestDatabase(IntegrationTest):
             "write_concern": WriteConcern(w=1),
             "read_concern": ReadConcern(level="local"),
         }
-        db2 = db1.with_options(**newopts)  # type: ignore[arg-type]
+        db2 = db1.with_options(**newopts)  # type: ignore[arg-type, call-overload]
         for opt in newopts:
             self.assertEqual(getattr(db2, opt), newopts.get(opt, getattr(db1, opt)))
 

test/test_typing.py

@@ -34,7 +34,7 @@ from typing import (
     cast,
 )
 
-try:
+if TYPE_CHECKING:
     from typing_extensions import NotRequired, TypedDict
 
     from bson import ObjectId
@@ -49,16 +49,13 @@ try:
         year: int
 
     class ImplicitMovie(TypedDict):
-        _id: NotRequired[ObjectId]  # pyright: ignore[reportGeneralTypeIssues]
+        _id: NotRequired[ObjectId]
         name: str
         year: int
-
-except ImportError:
-    Movie = dict  # type:ignore[misc,assignment]
-    ImplicitMovie = dict  # type: ignore[assignment,misc]
-    MovieWithId = dict  # type: ignore[assignment,misc]
-    TypedDict = None
-    NotRequired = None  # type: ignore[assignment]
+else:
+    Movie = dict
+    ImplicitMovie = dict
+    NotRequired = None
 
 
 try:
@@ -234,6 +231,19 @@ class TestPymongo(IntegrationTest):
                 execute_transaction, read_preference=ReadPreference.PRIMARY
             )
 
+    def test_with_options(self) -> None:
+        coll: Collection[Dict[str, Any]] = self.coll
+        coll.drop()
+        doc = {"name": "foo", "year": 1982, "other": 1}
+        coll.insert_one(doc)
+
+        coll2 = coll.with_options(codec_options=CodecOptions(document_class=Movie))
+        retrieved = coll2.find_one()
+        assert retrieved is not None
+        assert retrieved["name"] == "foo"
+        # We expect a type error here.
+        assert retrieved["other"] == 1  # type:ignore[typeddict-item]
+
 
 class TestDecode(unittest.TestCase):
     def test_bson_decode(self) -> None:
@@ -426,7 +436,7 @@ class TestDocumentType(PyMongoTestCase):
         )
         coll.bulk_write(
             [
-                InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971})
+                InsertOne({"_id": ObjectId(), "name": "THX-1138", "year": 1971})  # pyright: ignore
             ]  # No error because it is in-line.
         )
 
@@ -443,7 +453,7 @@ class TestDocumentType(PyMongoTestCase):
         )
         coll.bulk_write(
             [
-                ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971})
+                ReplaceOne({}, {"_id": ObjectId(), "name": "THX-1138", "year": 1971})  # pyright: ignore
             ]  # No error because it is in-line.
         )
 
@@ -566,7 +576,7 @@ class TestCodecOptionsDocumentType(unittest.TestCase):
     def test_typeddict_document_type(self) -> None:
         options: CodecOptions[Movie] = CodecOptions()
         # Suppress: Cannot instantiate type "Type[Movie]".
-        obj = options.document_class(name="a", year=1)  # type: ignore[misc]
+        obj = options.document_class(name="a", year=1)
         assert obj["year"] == 1
         assert obj["name"] == "a"
 

tools/synchro.py

@@ -23,7 +23,7 @@ import re
 from os import listdir
 from pathlib import Path
 
-from unasync import Rule, unasync_files  # type: ignore[import]
+from unasync import Rule, unasync_files  # type: ignore[import-not-found]
 
 replacements = {
     "AsyncCollection": "Collection",