PYTHON-3806 add types to message.py (#1312)

Author: sleepyStick

pymongo/bulk.py

@@ -153,7 +153,7 @@ def _merge_command(
         full_result["writeConcernErrors"].append(wce)
 
 
-def _raise_bulk_write_error(full_result: Mapping[str, Any]) -> NoReturn:
+def _raise_bulk_write_error(full_result: _DocumentOut) -> NoReturn:
     """Raise a BulkWriteError from the full bulk api result."""
     if full_result["writeErrors"]:
         full_result["writeErrors"].sort(key=lambda error: error["index"])

pymongo/client_session.py

@@ -986,7 +986,7 @@ def _apply_to(
         self,
         command: MutableMapping[str, Any],
         is_retryable: bool,
-        read_preference: ReadPreference,
+        read_preference: _ServerMode,
         conn: Connection,
     ) -> None:
         self._check_ended()

pymongo/command_cursor.py

@@ -33,7 +33,7 @@
 from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure
 from pymongo.message import _CursorAddress, _GetMore, _OpMsg, _OpReply, _RawBatchGetMore
 from pymongo.response import PinnedResponse
-from pymongo.typings import _Address, _DocumentType
+from pymongo.typings import _Address, _DocumentOut, _DocumentType
 
 if TYPE_CHECKING:
     from pymongo.client_session import ClientSession
@@ -94,6 +94,7 @@ def __die(self, synchronous: bool = False) -> None:
         self.__killed = True
         if self.__id and not already_killed:
             cursor_id = self.__id
+            assert self.__address is not None
             address = _CursorAddress(self.__address, self.__ns)
         else:
             # Skip killCursors.
@@ -219,7 +220,7 @@ def _unpack_response(
         codec_options: CodecOptions[Mapping[str, Any]],
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[Mapping[str, Any]]:
+    ) -> List[_DocumentOut]:
         return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response)
 
     def _refresh(self) -> int:
@@ -380,7 +381,7 @@ def __init__(
             comment,
         )
 
-    def _unpack_response(
+    def _unpack_response(  # type: ignore[override]
         self,
         response: Union[_OpReply, _OpMsg],
         cursor_id: Optional[int],
@@ -393,7 +394,7 @@ def _unpack_response(
             # OP_MSG returns firstBatch/nextBatch documents as a BSON array
             # Re-assemble the array of documents into a document stream
             _convert_raw_document_lists_to_streams(raw_response[0])
-        return raw_response
+        return raw_response  # type: ignore[return-value]
 
     def __getitem__(self, index: int) -> NoReturn:
         raise InvalidOperation("Cannot call __getitem__ on RawBatchCursor")

pymongo/cursor.py

@@ -57,7 +57,7 @@
     _RawBatchQuery,
 )
 from pymongo.response import PinnedResponse
-from pymongo.typings import _Address, _CollationIn, _DocumentType
+from pymongo.typings import _Address, _CollationIn, _DocumentOut, _DocumentType
 
 if TYPE_CHECKING:
     from _typeshed import SupportsItems
@@ -420,6 +420,7 @@ def __die(self, synchronous: bool = False) -> None:
         self.__killed = True
         if self.__id and not already_killed:
             cursor_id = self.__id
+            assert self.__address is not None
             address = _CursorAddress(self.__address, f"{self.__dbname}.{self.__collname}")
         else:
             # Skip killCursors.
@@ -1129,7 +1130,7 @@ def _unpack_response(
         codec_options: CodecOptions,
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[Mapping[str, Any]]:
+    ) -> List[_DocumentOut]:
         return response.unpack_response(cursor_id, codec_options, user_fields, legacy_response)
 
     def _read_preference(self) -> _ServerMode:
@@ -1355,13 +1356,13 @@ def _unpack_response(
         codec_options: CodecOptions[Mapping[str, Any]],
         user_fields: Optional[Mapping[str, Any]] = None,
         legacy_response: bool = False,
-    ) -> List[Mapping[str, Any]]:
+    ) -> List[_DocumentOut]:
         raw_response = response.raw_response(cursor_id, user_fields=user_fields)
         if not legacy_response:
             # OP_MSG returns firstBatch/nextBatch documents as a BSON array
             # Re-assemble the array of documents into a document stream
             _convert_raw_document_lists_to_streams(raw_response[0])
-        return raw_response
+        return cast(List["_DocumentOut"], raw_response)
 
     def explain(self) -> _DocumentType:
         """Returns an explain plan record for this cursor.

pymongo/encryption.py

@@ -80,8 +80,6 @@
 if TYPE_CHECKING:
     from pymongocrypt.mongocrypt import MongoCryptKmsContext
 
-    from pymongo.response import Response
-
 _HTTPS_PORT = 443
 _KMS_CONNECT_TIMEOUT = CONNECT_TIMEOUT  # CDRIVER-3262 redefined this value to CONNECT_TIMEOUT
 _MONGOCRYPTD_TIMEOUT_MS = 10000
@@ -409,7 +407,7 @@ def encrypt(
             encrypt_cmd = _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
             return encrypt_cmd
 
-    def decrypt(self, response: Response) -> Optional[bytes]:
+    def decrypt(self, response: bytes) -> Optional[bytes]:
         """Decrypt a MongoDB command response.
 
         :Parameters:

pymongo/errors.py

@@ -13,10 +13,25 @@
 # limitations under the License.
 
 """Exceptions raised by PyMongo."""
-from typing import Any, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 from bson.errors import InvalidDocument
 
+if TYPE_CHECKING:
+    from pymongo.typings import _DocumentOut
+
 try:
     # CPython 3.7+
     from ssl import SSLCertVerificationError as _CertificateError
@@ -286,9 +301,9 @@ class BulkWriteError(OperationFailure):
     .. versionadded:: 2.7
     """
 
-    details: Mapping[str, Any]
+    details: _DocumentOut
 
-    def __init__(self, results: Mapping[str, Any]) -> None:
+    def __init__(self, results: _DocumentOut) -> None:
         super().__init__("batch op errors occurred", 65, results)
 
     def __reduce__(self) -> Tuple[Any, Any]:

pymongo/helpers.py

@@ -53,6 +53,7 @@
 if TYPE_CHECKING:
     from pymongo.cursor import _Hint
     from pymongo.operations import _IndexList
+    from pymongo.typings import _DocumentOut
 
 # From the SDAM spec, the "node is shutting down" codes.
 _SHUTDOWN_CODES: frozenset = frozenset(
@@ -156,7 +157,7 @@ def _index_document(index_list: _IndexList) -> SON[str, Any]:
 
 
 def _check_command_response(
-    response: Mapping[str, Any],
+    response: _DocumentOut,
     max_wire_version: Optional[int],
     allowable_errors: Optional[Container[Union[int, str]]] = None,
     parse_write_concern_error: bool = False,