PYTHON-3813 add types to pool.py (#1318)

Author: sleepyStick

pymongo/auth.py

@@ -239,12 +239,15 @@ def _authenticate_scram(credentials: MongoCredential, conn: Connection, mechanis
 
     ctx = conn.auth_ctx
     if ctx and ctx.speculate_succeeded():
+        assert isinstance(ctx, _ScramContext)
+        assert ctx.scram_data is not None
         nonce, first_bare = ctx.scram_data
         res = ctx.speculative_authenticate
     else:
         nonce, first_bare, cmd = _authenticate_scram_start(credentials, mechanism)
         res = conn.command(source, cmd)
 
+    assert res is not None
     server_first = res["payload"]
     parsed = _parse_scram_response(server_first)
     iterations = int(parsed[b"i"])
@@ -575,7 +578,7 @@ def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
 
 
 class _X509Context(_AuthContext):
-    def speculate_command(self) -> Optional[MutableMapping[str, Any]]:
+    def speculate_command(self) -> MutableMapping[str, Any]:
         cmd = SON([("authenticate", 1), ("mechanism", "MONGODB-X509")])
         if self.credentials.username is not None:
             cmd["user"] = self.credentials.username

pymongo/auth_oidc.py

@@ -19,7 +19,17 @@
 import threading
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+)
 
 import bson
 from bson.binary import Binary
@@ -242,7 +252,9 @@ def clear(self) -> None:
         self.idp_resp = None
         self.token_exp_utc = None
 
-    def run_command(self, conn: Connection, cmd: Mapping[str, Any]) -> Optional[Mapping[str, Any]]:
+    def run_command(
+        self, conn: Connection, cmd: MutableMapping[str, Any]
+    ) -> Optional[Mapping[str, Any]]:
         try:
             return conn.command("$external", cmd, no_reauth=True)  # type: ignore[call-arg]
         except OperationFailure as exc:
@@ -276,6 +288,7 @@ def authenticate(
             assert cmd is not None
             resp = self.run_command(conn, cmd)
 
+        assert resp is not None
         if resp["done"]:
             conn.oidc_token_gen_id = self.token_gen_id
             return None
@@ -297,6 +310,7 @@ def authenticate(
             ]
         )
         resp = self.run_command(conn, cmd)
+        assert resp is not None
         if not resp["done"]:
             self.clear()
             raise OperationFailure("SASL conversation failed to complete.")

pymongo/client_options.py

@@ -15,7 +15,7 @@
 """Tools to parse mongo client options."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence, Tuple, cast
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sequence, Tuple, cast
 
 from bson.codec_options import _parse_codec_options
 from pymongo import common
@@ -309,11 +309,12 @@ def load_balanced(self) -> Optional[bool]:
         return self.__load_balanced
 
     @property
-    def event_listeners(self) -> _EventListeners:
+    def event_listeners(self) -> List[_EventListeners]:
         """The event listeners registered for this client.
 
         See :mod:`~pymongo.monitoring` for details.
 
         .. versionadded:: 4.0
         """
+        assert self.__pool_options._event_listeners is not None
         return self.__pool_options._event_listeners.event_listeners()

pymongo/collection.py

@@ -20,7 +20,6 @@
     TYPE_CHECKING,
     Any,
     Callable,
-    Container,
     ContextManager,
     Generic,
     Iterable,
@@ -268,11 +267,11 @@ def _conn_for_writes(self, session: Optional[ClientSession]) -> ContextManager[C
     def _command(
         self,
         conn: Connection,
-        command: Mapping[str, Any],
+        command: MutableMapping[str, Any],
         read_preference: Optional[_ServerMode] = None,
         codec_options: Optional[CodecOptions] = None,
         check: bool = True,
-        allowable_errors: Optional[Container[Any]] = None,
+        allowable_errors: Optional[Sequence[Union[str, int]]] = None,
         read_concern: Optional[ReadConcern] = None,
         write_concern: Optional[WriteConcern] = None,
         collation: Optional[_CollationIn] = None,
@@ -1753,7 +1752,7 @@ def _count_cmd(
         session: Optional[ClientSession],
         conn: Connection,
         read_preference: Optional[_ServerMode],
-        cmd: Mapping[str, Any],
+        cmd: SON[str, Any],
         collation: Optional[Collation],
     ) -> int:
         """Internal count command helper."""
@@ -1777,7 +1776,7 @@ def _aggregate_one_result(
         self,
         conn: Connection,
         read_preference: Optional[_ServerMode],
-        cmd: Mapping[str, Any],
+        cmd: SON[str, Any],
         collation: Optional[_CollationIn],
         session: Optional[ClientSession],
     ) -> Optional[Mapping[str, Any]]:

pymongo/command_cursor.py

@@ -25,6 +25,7 @@
     Mapping,
     NoReturn,
     Optional,
+    Sequence,
     Union,
 )
 
@@ -220,7 +221,7 @@ def _unpack_response(
         codec_options: CodecOptions[Mapping[str, Any]],
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[_DocumentOut]:
+    ) -> Sequence[_DocumentOut]:
         return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response)
 
     def _refresh(self) -> int:

pymongo/compression_support.py

@@ -14,7 +14,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Iterable, List, Union
+from typing import Any, Iterable, List, Optional, Union
 
 try:
     import snappy
@@ -96,7 +96,7 @@ def __init__(self, compressors: List[str], zlib_compression_level: int):
         self.zlib_compression_level = zlib_compression_level
 
     def get_compression_context(
-        self, compressors: List[str]
+        self, compressors: Optional[List[str]]
     ) -> Union[SnappyContext, ZlibContext, ZstdContext, None]:
         if compressors:
             chosen = compressors[0]

pymongo/cursor.py

@@ -1130,7 +1130,7 @@ def _unpack_response(
         codec_options: CodecOptions,
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[_DocumentOut]:
+    ) -> Sequence[_DocumentOut]:
         return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response)
 
     def _read_preference(self) -> _ServerMode:

pymongo/database.py

@@ -694,7 +694,7 @@ def _command(
         value: int = 1,
         check: bool = True,
         allowable_errors: Optional[Sequence[Union[str, int]]] = None,
-        read_preference: Optional[_ServerMode] = ReadPreference.PRIMARY,
+        read_preference: _ServerMode = ReadPreference.PRIMARY,
         codec_options: CodecOptions[Dict[str, Any]] = DEFAULT_CODEC_OPTIONS,
         write_concern: Optional[WriteConcern] = None,
         parse_write_concern_error: bool = False,
@@ -711,7 +711,7 @@ def _command(
         value: int = 1,
         check: bool = True,
         allowable_errors: Optional[Sequence[Union[str, int]]] = None,
-        read_preference: Optional[_ServerMode] = ReadPreference.PRIMARY,
+        read_preference: _ServerMode = ReadPreference.PRIMARY,
         codec_options: CodecOptions[_CodecDocumentType] = ...,
         write_concern: Optional[WriteConcern] = None,
         parse_write_concern_error: bool = False,
@@ -727,7 +727,7 @@ def _command(
         value: int = 1,
         check: bool = True,
         allowable_errors: Optional[Sequence[Union[str, int]]] = None,
-        read_preference: Optional[_ServerMode] = ReadPreference.PRIMARY,
+        read_preference: _ServerMode = ReadPreference.PRIMARY,
         codec_options: Union[
             CodecOptions[Dict[str, Any]], CodecOptions[_CodecDocumentType]
         ] = DEFAULT_CODEC_OPTIONS,

pymongo/message.py

@@ -1018,7 +1018,7 @@ def unack_write(
             cmd = self._start(cmd, request_id, docs)
             start = datetime.datetime.now()
         try:
-            result = self.conn.unack_write(msg, max_doc_size)
+            result = self.conn.unack_write(msg, max_doc_size)  # type: ignore[func-returns-value]
             if self.publish:
                 duration = (datetime.datetime.now() - start) + duration
                 if result is not None:
@@ -1050,7 +1050,7 @@ def write_command(
         request_id: int,
         msg: bytes,
         docs: List[Mapping[str, Any]],
-    ) -> Mapping[str, Any]:
+    ) -> Dict[str, Any]:
         """A proxy for SocketInfo.write_command that handles event publishing."""
         if self.publish:
             assert self.start_time is not None
@@ -1127,7 +1127,7 @@ class _EncryptedBulkWriteContext(_BulkWriteContext):
 
     def __batch_command(
         self, cmd: MutableMapping[str, Any], docs: List[Mapping[str, Any]]
-    ) -> Tuple[Mapping[str, Any], List[Mapping[str, Any]]]:
+    ) -> Tuple[Dict[str, Any], List[Mapping[str, Any]]]:
         namespace = self.db_name + ".$cmd"
         msg, to_send = _encode_batched_write_command(
             namespace, self.op_type, cmd, docs, self.codec, self
@@ -1517,7 +1517,7 @@ def unpack_response(
         codec_options: CodecOptions = _UNICODE_REPLACE_CODEC_OPTIONS,
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[_DocumentOut]:
+    ) -> List[Dict[str, Any]]:
         """Unpack a response from the database and decode the BSON document(s).
 
         Check the response for errors and unpack, returning a dictionary
@@ -1541,7 +1541,7 @@ def unpack_response(
             return bson.decode_all(self.documents, codec_options)
         return bson._decode_all_selective(self.documents, codec_options, user_fields)
 
-    def command_response(self, codec_options: CodecOptions) -> Mapping[str, Any]:
+    def command_response(self, codec_options: CodecOptions) -> Dict[str, Any]:
         """Unpack a command response."""
         docs = self.unpack_response(codec_options=codec_options)
         assert self.number_returned == 1
@@ -1604,7 +1604,7 @@ def unpack_response(
         codec_options: CodecOptions = _UNICODE_REPLACE_CODEC_OPTIONS,
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[_DocumentOut]:
+    ) -> List[Dict[str, Any]]:
         """Unpack a OP_MSG command response.
 
         :Parameters:
@@ -1619,7 +1619,7 @@ def unpack_response(
         assert not legacy_response
         return bson._decode_all_selective(self.payload_document, codec_options, user_fields)
 
-    def command_response(self, codec_options: CodecOptions) -> Mapping[str, Any]:
+    def command_response(self, codec_options: CodecOptions) -> Dict[str, Any]:
         """Unpack a command response."""
         return self.unpack_response(codec_options=codec_options)[0]
 

pymongo/mongo_client.py

@@ -116,6 +116,7 @@
     import sys
     from types import TracebackType
 
+    from bson.objectid import ObjectId
     from pymongo.bulk import _Bulk
     from pymongo.client_session import ClientSession, _ServerSession
     from pymongo.cursor import _ConnectionManager
@@ -1898,7 +1899,9 @@ def _tmp_session(
         else:
             yield None
 
-    def _send_cluster_time(self, command: MutableMapping[str, Any], session: ClientSession) -> None:
+    def _send_cluster_time(
+        self, command: MutableMapping[str, Any], session: Optional[ClientSession]
+    ) -> None:
         topology_time = self._topology.max_cluster_time()
         session_time = session.cluster_time if session else None
         if topology_time and session_time:
@@ -2255,7 +2258,7 @@ def __init__(self, client: MongoClient, server: Server, session: Optional[Client
         # of the pool at the time the connection attempt was started."
         self.sock_generation = server.pool.gen.get_overall()
         self.completed_handshake = False
-        self.service_id = None
+        self.service_id: Optional[ObjectId] = None
         self.handled = False
 
     def contribute_socket(self, conn: Connection, completed_handshake: bool = True) -> None:

pymongo/monitor.py

@@ -32,7 +32,7 @@
 from pymongo.srv_resolver import _SrvResolver
 
 if TYPE_CHECKING:
-    from pymongo.pool import Connection, Pool
+    from pymongo.pool import Connection, Pool, _CancellationContext
     from pymongo.settings import TopologySettings
     from pymongo.topology import Topology
 
@@ -131,9 +131,8 @@ def __init__(
         self._pool = pool
         self._settings = topology_settings
         self._listeners = self._settings._pool_options._event_listeners
-        pub = self._listeners is not None
-        self._publish = pub and self._listeners.enabled_for_server_heartbeat
-        self._cancel_context = None
+        self._publish = self._listeners is not None and self._listeners.enabled_for_server_heartbeat
+        self._cancel_context: Optional[_CancellationContext] = None
         self._rtt_monitor = _RttMonitor(
             topology,
             topology_settings,
@@ -238,7 +237,8 @@ def _check_server(self) -> ServerDescription:
             address = sd.address
             duration = time.monotonic() - start
             if self._publish:
-                awaited = sd.is_server_type_known and sd.topology_version
+                awaited = bool(sd.is_server_type_known and sd.topology_version)
+                assert self._listeners is not None
                 self._listeners.publish_server_heartbeat_failed(address, duration, error, awaited)
             self._reset_connection()
             if isinstance(error, _OperationCancelled):
@@ -254,6 +254,7 @@ def _check_once(self) -> ServerDescription:
         """
         address = self._server_description.address
         if self._publish:
+            assert self._listeners is not None
             self._listeners.publish_server_heartbeat_started(address)
 
         if self._cancel_context and self._cancel_context.cancelled:
@@ -267,6 +268,7 @@ def _check_once(self) -> ServerDescription:
             avg_rtt, min_rtt = self._rtt_monitor.get()
             sd = ServerDescription(address, response, avg_rtt, min_round_trip_time=min_rtt)
             if self._publish:
+                assert self._listeners is not None
                 self._listeners.publish_server_heartbeat_succeeded(
                     address, round_trip_time, response, response.awaitable
                 )

pymongo/network.py

@@ -47,14 +47,13 @@
 if TYPE_CHECKING:
     from bson import CodecOptions
     from pymongo.client_session import ClientSession
-    from pymongo.collation import Collation
     from pymongo.compression_support import SnappyContext, ZlibContext, ZstdContext
     from pymongo.mongo_client import MongoClient
     from pymongo.monitoring import _EventListeners
     from pymongo.pool import Connection
     from pymongo.read_concern import ReadConcern
     from pymongo.read_preferences import _ServerMode
-    from pymongo.typings import _Address, _DocumentOut, _DocumentType
+    from pymongo.typings import _Address, _CollationIn, _DocumentOut, _DocumentType
     from pymongo.write_concern import WriteConcern
 
 _UNPACK_HEADER = struct.Struct("<iiii").unpack
@@ -65,7 +64,7 @@ def command(
     dbname: str,
     spec: MutableMapping[str, Any],
     is_mongos: bool,
-    read_preference: _ServerMode,
+    read_preference: Optional[_ServerMode],
     codec_options: CodecOptions[_DocumentType],
     session: Optional[ClientSession],
     client: Optional[MongoClient],
@@ -76,7 +75,7 @@ def command(
     max_bson_size: Optional[int] = None,
     read_concern: Optional[ReadConcern] = None,
     parse_write_concern_error: bool = False,
-    collation: Optional[Collation] = None,
+    collation: Optional[_CollationIn] = None,
     compression_ctx: Union[SnappyContext, ZlibContext, ZstdContext, None] = None,
     use_op_msg: bool = False,
     unacknowledged: bool = False,
@@ -119,6 +118,7 @@ def command(
     # Publish the original command document, perhaps with lsid and $clusterTime.
     orig = spec
     if is_mongos and not use_op_msg:
+        assert read_preference is not None
         spec = message._maybe_add_read_preference(spec, read_preference)
     if read_concern and not (session and session.in_transaction):
         if read_concern.level:
@@ -232,7 +232,7 @@ def command(
 
 
 def receive_message(
-    conn: Connection, request_id: int, max_message_size: int = MAX_MESSAGE_SIZE
+    conn: Connection, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
     if _csot.get_timeout():