PYTHON-3060 Add typings to pymongo package (mongodb#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