PYTHON-3060 Add typings to pymongo package (#831)

Author: blink1073

bson/__init__.py

@@ -61,10 +61,10 @@
 import struct
 import sys
 import uuid
-from codecs import utf_8_decode as _utf_8_decode  # type: ignore
-from codecs import utf_8_encode as _utf_8_encode  # type: ignore
+from codecs import utf_8_decode as _utf_8_decode  # type: ignore[attr-defined]
+from codecs import utf_8_encode as _utf_8_encode  # type: ignore[attr-defined]
 from collections import abc as _abc
-from typing import (TYPE_CHECKING, Any, BinaryIO, Callable, Dict, Generator,
+from typing import (IO, TYPE_CHECKING, Any, BinaryIO, Callable, Dict, Generator,
                     Iterator, List, Mapping, MutableMapping, NoReturn,
                     Sequence, Tuple, Type, TypeVar, Union, cast)
 
@@ -88,11 +88,13 @@
 
 # Import RawBSONDocument for type-checking only to avoid circular dependency.
 if TYPE_CHECKING:
+    from array import array
+    from mmap import mmap
     from bson.raw_bson import RawBSONDocument
 
 
 try:
-    from bson import _cbson  # type: ignore
+    from bson import _cbson  # type: ignore[attr-defined]
     _USE_C = True
 except ImportError:
     _USE_C = False
@@ -851,6 +853,7 @@ def _datetime_to_millis(dtm: datetime.datetime) -> int:
 
 _DocumentIn = Mapping[str, Any]
 _DocumentOut = Union[MutableMapping[str, Any], "RawBSONDocument"]
+_ReadableBuffer = Union[bytes, memoryview, "mmap", "array"]
 
 
 def encode(document: _DocumentIn, check_keys: bool = False, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> bytes:
@@ -880,7 +883,7 @@ def encode(document: _DocumentIn, check_keys: bool = False, codec_options: Codec
     return _dict_to_bson(document, check_keys, codec_options)
 
 
-def decode(data: bytes, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> _DocumentOut:
+def decode(data: _ReadableBuffer, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> Dict[str, Any]:
     """Decode BSON to a document.
 
     By default, returns a BSON document represented as a Python
@@ -912,7 +915,7 @@ def decode(data: bytes, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) ->
     return _bson_to_dict(data, codec_options)
 
 
-def decode_all(data: bytes, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> List[_DocumentOut]:
+def decode_all(data: _ReadableBuffer, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> List[Dict[str, Any]]:
     """Decode BSON data to multiple documents.
 
     `data` must be a bytes-like object implementing the buffer protocol that
@@ -1075,7 +1078,7 @@ def decode_iter(data: bytes, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS
         yield _bson_to_dict(elements, codec_options)
 
 
-def decode_file_iter(file_obj: BinaryIO, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> Iterator[_DocumentOut]:
+def decode_file_iter(file_obj: Union[BinaryIO, IO], codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> Iterator[_DocumentOut]:
     """Decode bson data from a file to multiple documents as a generator.
 
     Works similarly to the decode_all function, but reads from the file object
@@ -1158,7 +1161,7 @@ def encode(cls: Type["BSON"], document: _DocumentIn, check_keys: bool = False,
         """
         return cls(encode(document, check_keys, codec_options))
 
-    def decode(self, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> _DocumentOut:  # type: ignore[override]
+    def decode(self, codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS) -> Dict[str, Any]:  # type: ignore[override]
         """Decode this BSON data.
 
         By default, returns a BSON document represented as a Python

bson/binary.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Tuple, Type
+from typing import Any, Tuple, Type, Union, TYPE_CHECKING
 from uuid import UUID
 
 """Tools for representing BSON binary data.
@@ -57,6 +57,11 @@
 """
 
 
+if TYPE_CHECKING:
+    from array import array as _array
+    from mmap import mmap as _mmap
+
+
 class UuidRepresentation:
     UNSPECIFIED = 0
     """An unspecified UUID representation.
@@ -211,7 +216,7 @@ class Binary(bytes):
     _type_marker = 5
     __subtype: int
 
-    def __new__(cls: Type["Binary"], data: bytes, subtype: int = BINARY_SUBTYPE) -> "Binary":
+    def __new__(cls: Type["Binary"], data: Union[memoryview, bytes, "_mmap", "_array"], subtype: int = BINARY_SUBTYPE) -> "Binary":
         if not isinstance(subtype, int):
             raise TypeError("subtype must be an instance of int")
         if subtype >= 256 or subtype < 0:

gridfs/grid_file.py

@@ -874,10 +874,10 @@ def next(self) -> GridOut:
 
     __next__ = next
 
-    def add_option(self, *args: Any, **kwargs: Any) -> None:
+    def add_option(self, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
         raise NotImplementedError("Method does not exist for GridOutCursor")
 
-    def remove_option(self, *args: Any, **kwargs: Any) -> None:
+    def remove_option(self, *args: Any, **kwargs: Any) -> None:  # type: ignore[override]
         raise NotImplementedError("Method does not exist for GridOutCursor")
 
     def _clone_base(self, session: ClientSession) -> "GridOutCursor":

mypy.ini

@@ -1,11 +1,33 @@
 [mypy]
+check_untyped_defs = true
 disallow_subclassing_any = true
 disallow_incomplete_defs = true
 no_implicit_optional = true
+pretty = true
+show_error_context = true
+show_error_codes = true
 strict_equality = true
 warn_unused_configs = true
 warn_unused_ignores = true
 warn_redundant_casts = true
 
+[mypy-kerberos.*]
+ignore_missing_imports = True
+
 [mypy-mockupdb]
 ignore_missing_imports = True
+
+[mypy-pymongo_auth_aws.*]
+ignore_missing_imports = True
+
+[mypy-pymongocrypt.*]
+ignore_missing_imports = True
+
+[mypy-service_identity.*]
+ignore_missing_imports = True
+
+[mypy-snappy.*]
+ignore_missing_imports = True
+
+[mypy-winkerberos.*]
+ignore_missing_imports = True

pymongo/__init__.py

@@ -14,6 +14,8 @@
 
 """Python driver for MongoDB."""
 
+from typing import Tuple, Union
+
 ASCENDING = 1
 """Ascending sort order."""
 DESCENDING = -1
@@ -53,35 +55,33 @@
 .. _text index: http://docs.mongodb.org/manual/core/index-text/
 """
 
-version_tuple = (4, 1, 0, '.dev0')
+version_tuple: Tuple[Union[int, str], ...] = (4, 1, 0, '.dev0')
 
-def get_version_string():
+def get_version_string() -> str:
     if isinstance(version_tuple[-1], str):
         return '.'.join(map(str, version_tuple[:-1])) + version_tuple[-1]
     return '.'.join(map(str, version_tuple))
 
-__version__ = version = get_version_string()
+__version__: str = get_version_string()
+version = __version__
+
 """Current version of PyMongo."""
 
 from pymongo.collection import ReturnDocument
-from pymongo.common import (MIN_SUPPORTED_WIRE_VERSION,
-                            MAX_SUPPORTED_WIRE_VERSION)
+from pymongo.common import (MAX_SUPPORTED_WIRE_VERSION,
+                            MIN_SUPPORTED_WIRE_VERSION)
 from pymongo.cursor import CursorType
 from pymongo.mongo_client import MongoClient
-from pymongo.operations import (IndexModel,
-                                InsertOne,
-                                DeleteOne,
-                                DeleteMany,
-                                UpdateOne,
-                                UpdateMany,
-                                ReplaceOne)
+from pymongo.operations import (DeleteMany, DeleteOne, IndexModel, InsertOne,
+                                ReplaceOne, UpdateMany, UpdateOne)
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
-def has_c():
+
+def has_c() -> bool:
     """Is the C extension installed?"""
     try:
-        from pymongo import _cmessage
+        from pymongo import _cmessage  # type: ignore[attr-defined]
         return True
     except ImportError:
         return False

pymongo/aggregation.py

@@ -15,11 +15,10 @@
 """Perform aggregation operations on a collection or database."""
 
 from bson.son import SON
-
 from pymongo import common
 from pymongo.collation import validate_collation_or_none
 from pymongo.errors import ConfigurationError
-from pymongo.read_preferences import _AggWritePref, ReadPreference
+from pymongo.read_preferences import ReadPreference, _AggWritePref
 
 
 class _AggregationCommand(object):
@@ -37,7 +36,7 @@ def __init__(self, target, cursor_class, pipeline, options,
 
         self._target = target
 
-        common.validate_list('pipeline', pipeline)
+        pipeline = common.validate_list('pipeline', pipeline)
         self._pipeline = pipeline
         self._performs_write = False
         if pipeline and ("$out" in pipeline[-1] or "$merge" in pipeline[-1]):
@@ -82,7 +81,6 @@ def _cursor_namespace(self):
         """The namespace in which the aggregate command is run."""
         raise NotImplementedError
 
-    @property
     def _cursor_collection(self, cursor_doc):
         """The Collection used for the aggregate command cursor."""
         raise NotImplementedError

pymongo/auth.py

@@ -19,9 +19,9 @@
 import hmac
 import os
 import socket
-
 from base64 import standard_b64decode, standard_b64encode
 from collections import namedtuple
+from typing import Callable, Mapping
 from urllib.parse import quote
 
 from bson.binary import Binary
@@ -97,7 +97,7 @@ def __hash__(self):
 """Mechanism properties for GSSAPI authentication."""
 
 
-_AWSProperties = namedtuple('AWSProperties', ['aws_session_token'])
+_AWSProperties = namedtuple('_AWSProperties', ['aws_session_token'])
 """Mechanism properties for MONGODB-AWS authentication."""
 
 
@@ -140,9 +140,9 @@ def _build_credentials_tuple(mech, source, user, passwd, extra, database):
 
         properties = extra.get('authmechanismproperties', {})
         aws_session_token = properties.get('AWS_SESSION_TOKEN')
-        props = _AWSProperties(aws_session_token=aws_session_token)
+        aws_props = _AWSProperties(aws_session_token=aws_session_token)
         # user can be None for temporary link-local EC2 credentials.
-        return MongoCredential(mech, '$external', user, passwd, props, None)
+        return MongoCredential(mech, '$external', user, passwd, aws_props, None)
     elif mech == 'PLAIN':
         source_database = source or database or '$external'
         return MongoCredential(mech, source_database, user, passwd, None, None)
@@ -471,7 +471,7 @@ def _authenticate_default(credentials, sock_info):
         return _authenticate_scram(credentials, sock_info, 'SCRAM-SHA-1')
 
 
-_AUTH_MAP = {
+_AUTH_MAP: Mapping[str, Callable] = {
     'GSSAPI': _authenticate_gssapi,
     'MONGODB-CR': _authenticate_mongo_cr,
     'MONGODB-X509': _authenticate_x509,
@@ -532,7 +532,7 @@ def speculate_command(self):
         return cmd
 
 
-_SPECULATIVE_AUTH_MAP = {
+_SPECULATIVE_AUTH_MAP: Mapping[str, Callable] = {
     'MONGODB-X509': _X509Context,
     'SCRAM-SHA-1': functools.partial(_ScramContext, mechanism='SCRAM-SHA-1'),
     'SCRAM-SHA-256': functools.partial(_ScramContext,
@@ -544,6 +544,6 @@ def speculate_command(self):
 def authenticate(credentials, sock_info):
     """Authenticate sock_info."""
     mechanism = credentials.mechanism
-    auth_func = _AUTH_MAP.get(mechanism)
+    auth_func = _AUTH_MAP[mechanism]
     auth_func(credentials, sock_info)
 

pymongo/auth_aws.py

@@ -16,12 +16,11 @@
 
 try:
     import pymongo_auth_aws
-    from pymongo_auth_aws import (AwsCredential,
-                                  AwsSaslContext,
+    from pymongo_auth_aws import (AwsCredential, AwsSaslContext,
                                   PyMongoAuthAwsError)
     _HAVE_MONGODB_AWS = True
 except ImportError:
-    class AwsSaslContext(object):
+    class AwsSaslContext(object):  # type: ignore
         def __init__(self, credentials):
             pass
     _HAVE_MONGODB_AWS = False
@@ -32,7 +31,7 @@ def __init__(self, credentials):
 from pymongo.errors import ConfigurationError, OperationFailure
 
 
-class _AwsSaslContext(AwsSaslContext):
+class _AwsSaslContext(AwsSaslContext):  # type: ignore
     # Dependency injection:
     def binary_type(self):
         """Return the bson.binary.Binary type."""

pymongo/bulk.py

@@ -17,31 +17,23 @@
 .. versionadded:: 2.7
 """
 import copy
-
 from itertools import islice
 
 from bson.objectid import ObjectId
 from bson.raw_bson import RawBSONDocument
 from bson.son import SON
 from pymongo.client_session import _validate_session_write_concern
-from pymongo.common import (validate_is_mapping,
-                            validate_is_document_type,
-                            validate_ok_for_replace,
-                            validate_ok_for_update)
-from pymongo.helpers import _RETRYABLE_ERROR_CODES, _get_wce_doc
 from pymongo.collation import validate_collation_or_none
-from pymongo.errors import (BulkWriteError,
-                            ConfigurationError,
-                            InvalidOperation,
-                            OperationFailure)
-from pymongo.message import (_INSERT, _UPDATE, _DELETE,
-                             _randint,
-                             _BulkWriteContext,
-                             _EncryptedBulkWriteContext)
+from pymongo.common import (validate_is_document_type, validate_is_mapping,
+                            validate_ok_for_replace, validate_ok_for_update)
+from pymongo.errors import (BulkWriteError, ConfigurationError,
+                            InvalidOperation, OperationFailure)
+from pymongo.helpers import _RETRYABLE_ERROR_CODES, _get_wce_doc
+from pymongo.message import (_DELETE, _INSERT, _UPDATE, _BulkWriteContext,
+                             _EncryptedBulkWriteContext, _randint)
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
-
 _DELETE_ALL = 0
 _DELETE_ONE = 1