Skip to content

Commit a3720d9

Browse files
authored
PYTHON-3568 Intellisense highlights multiple PyMongo methods because of CodecOptions (#1139)
1 parent 2b21e73 commit a3720d9

14 files changed

+150
-66
lines changed

.github/workflows/test-python.yml

+12-1
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,18 @@ jobs:
6767
mypy --install-types --non-interactive bson/codec_options.py
6868
mypy --install-types --non-interactive --disable-error-code var-annotated --disable-error-code attr-defined --disable-error-code union-attr --disable-error-code assignment --disable-error-code no-redef --disable-error-code index --allow-redefinition --allow-untyped-globals --exclude "test/mypy_fails/*.*" test
6969
python -m pip install -U typing_extensions
70-
mypy --install-types --non-interactive test/test_mypy.py
70+
mypy --install-types --non-interactive test/test_typing.py test/test_typing_strict.py
71+
- name: Run mypy strict
72+
run: |
73+
mypy --strict test/test_typing_strict.py
74+
- name: Run pyright
75+
run: |
76+
python -m pip install -U pip pyright==1.1.290
77+
pyright test/test_typing.py test/test_typing_strict.py
78+
- name: Run pyright strict
79+
run: |
80+
echo '{"strict": ["tests/test_typing_strict.py"]}' >> pyrightconfig.json
81+
pyright test/test_typing_strict.py
7182
7283
linkcheck:
7384
name: Check Links

bson/__init__.py

+15-19
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,6 @@
101101
DEFAULT_CODEC_OPTIONS,
102102
CodecOptions,
103103
DatetimeConversion,
104-
_DocumentType,
105104
_raw_document_class,
106105
)
107106
from bson.datetime_ms import (
@@ -125,8 +124,7 @@
125124

126125
# Import some modules for type-checking only.
127126
if TYPE_CHECKING:
128-
from array import array
129-
from mmap import mmap
127+
from bson.typings import _DocumentIn, _DocumentType, _ReadableBuffer
130128

131129
try:
132130
from bson import _cbson # type: ignore[attr-defined]
@@ -986,12 +984,8 @@ def _dict_to_bson(doc: Any, check_keys: bool, opts: CodecOptions, top_level: boo
986984
_CODEC_OPTIONS_TYPE_ERROR = TypeError("codec_options must be an instance of CodecOptions")
987985

988986

989-
_DocumentIn = Mapping[str, Any]
990-
_ReadableBuffer = Union[bytes, memoryview, "mmap", "array"]
991-
992-
993987
def encode(
994-
document: _DocumentIn,
988+
document: "_DocumentIn",
995989
check_keys: bool = False,
996990
codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS,
997991
) -> bytes:
@@ -1022,8 +1016,8 @@ def encode(
10221016

10231017

10241018
def decode(
1025-
data: _ReadableBuffer, codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1026-
) -> _DocumentType:
1019+
data: "_ReadableBuffer", codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1020+
) -> "_DocumentType":
10271021
"""Decode BSON to a document.
10281022
10291023
By default, returns a BSON document represented as a Python
@@ -1056,11 +1050,13 @@ def decode(
10561050
return _bson_to_dict(data, opts)
10571051

10581052

1059-
def _decode_all(data: _ReadableBuffer, opts: "CodecOptions[_DocumentType]") -> List[_DocumentType]:
1053+
def _decode_all(
1054+
data: "_ReadableBuffer", opts: "CodecOptions[_DocumentType]"
1055+
) -> "List[_DocumentType]":
10601056
"""Decode a BSON data to multiple documents."""
10611057
data, view = get_data_and_view(data)
10621058
data_len = len(data)
1063-
docs: List[_DocumentType] = []
1059+
docs: "List[_DocumentType]" = []
10641060
position = 0
10651061
end = data_len - 1
10661062
use_raw = _raw_document_class(opts.document_class)
@@ -1091,8 +1087,8 @@ def _decode_all(data: _ReadableBuffer, opts: "CodecOptions[_DocumentType]") -> L
10911087

10921088

10931089
def decode_all(
1094-
data: _ReadableBuffer, codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1095-
) -> List[_DocumentType]:
1090+
data: "_ReadableBuffer", codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1091+
) -> "List[_DocumentType]":
10961092
"""Decode BSON data to multiple documents.
10971093
10981094
`data` must be a bytes-like object implementing the buffer protocol that
@@ -1213,7 +1209,7 @@ def _decode_all_selective(data: Any, codec_options: CodecOptions, fields: Any) -
12131209
# Decode documents for internal use.
12141210
from bson.raw_bson import RawBSONDocument
12151211

1216-
internal_codec_options = codec_options.with_options(
1212+
internal_codec_options: CodecOptions[RawBSONDocument] = codec_options.with_options(
12171213
document_class=RawBSONDocument, type_registry=None
12181214
)
12191215
_doc = _bson_to_dict(data, internal_codec_options)
@@ -1228,7 +1224,7 @@ def _decode_all_selective(data: Any, codec_options: CodecOptions, fields: Any) -
12281224

12291225
def decode_iter(
12301226
data: bytes, codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1231-
) -> Iterator[_DocumentType]:
1227+
) -> "Iterator[_DocumentType]":
12321228
"""Decode BSON data to multiple documents as a generator.
12331229
12341230
Works similarly to the decode_all function, but yields one document at a
@@ -1264,7 +1260,7 @@ def decode_iter(
12641260

12651261
def decode_file_iter(
12661262
file_obj: Union[BinaryIO, IO], codec_options: "Optional[CodecOptions[_DocumentType]]" = None
1267-
) -> Iterator[_DocumentType]:
1263+
) -> "Iterator[_DocumentType]":
12681264
"""Decode bson data from a file to multiple documents as a generator.
12691265
12701266
Works similarly to the decode_all function, but reads from the file object
@@ -1325,7 +1321,7 @@ class BSON(bytes):
13251321
@classmethod
13261322
def encode(
13271323
cls: Type["BSON"],
1328-
document: _DocumentIn,
1324+
document: "_DocumentIn",
13291325
check_keys: bool = False,
13301326
codec_options: CodecOptions = DEFAULT_CODEC_OPTIONS,
13311327
) -> "BSON":
@@ -1352,7 +1348,7 @@ def encode(
13521348
"""
13531349
return cls(encode(document, check_keys, codec_options))
13541350

1355-
def decode(self, codec_options: "CodecOptions[_DocumentType]" = DEFAULT_CODEC_OPTIONS) -> _DocumentType: # type: ignore[override,assignment]
1351+
def decode(self, codec_options: "CodecOptions[_DocumentType]" = DEFAULT_CODEC_OPTIONS) -> "_DocumentType": # type: ignore[override,assignment]
13561352
"""Decode this BSON data.
13571353
13581354
By default, returns a BSON document represented as a Python

bson/codec_options.pyi

+4-6
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ you get the error: "TypeError: 'type' object is not subscriptable".
2222
import datetime
2323
import abc
2424
import enum
25-
from typing import Tuple, Generic, Optional, Mapping, Any, TypeVar, Type, Dict, Iterable, Tuple, MutableMapping, Callable, Union
25+
from typing import Tuple, Generic, Optional, Mapping, Any, Type, Dict, Iterable, Tuple, Callable, Union
26+
from bson.typings import _DocumentType, _DocumentTypeArg
2627

2728

2829
class TypeEncoder(abc.ABC, metaclass=abc.ABCMeta):
@@ -52,9 +53,6 @@ class TypeRegistry:
5253
def __init__(self, type_codecs: Optional[Iterable[Codec]] = ..., fallback_encoder: Optional[Fallback] = ...) -> None: ...
5354
def __eq__(self, other: Any) -> Any: ...
5455

55-
56-
_DocumentType = TypeVar("_DocumentType", bound=Mapping[str, Any])
57-
5856
class DatetimeConversion(int, enum.Enum):
5957
DATETIME = ...
6058
DATETIME_CLAMP = ...
@@ -82,7 +80,7 @@ class CodecOptions(Tuple, Generic[_DocumentType]):
8280
) -> CodecOptions[_DocumentType]: ...
8381

8482
# CodecOptions API
85-
def with_options(self, **kwargs: Any) -> CodecOptions[_DocumentType]: ...
83+
def with_options(self, **kwargs: Any) -> CodecOptions[_DocumentTypeArg]: ...
8684

8785
def _arguments_repr(self) -> str: ...
8886

@@ -100,7 +98,7 @@ class CodecOptions(Tuple, Generic[_DocumentType]):
10098
_fields: Tuple[str]
10199

102100

103-
DEFAULT_CODEC_OPTIONS: CodecOptions[MutableMapping[str, Any]]
101+
DEFAULT_CODEC_OPTIONS: "CodecOptions[Mapping[str, Any]]"
104102
_RAW_BSON_DOCUMENT_MARKER: int
105103

106104
def _raw_document_class(document_class: Any) -> bool: ...

bson/typings.py

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# Copyright 2023-Present MongoDB, Inc.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Type aliases used by bson"""
16+
from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, TypeVar, Union
17+
18+
if TYPE_CHECKING:
19+
from array import array
20+
from mmap import mmap
21+
22+
from bson.raw_bson import RawBSONDocument
23+
24+
25+
# Common Shared Types.
26+
_DocumentIn = Union[MutableMapping[str, Any], "RawBSONDocument"]
27+
_DocumentOut = _DocumentIn
28+
_DocumentType = TypeVar("_DocumentType", bound=Mapping[str, Any])
29+
_DocumentTypeArg = TypeVar("_DocumentTypeArg", bound=Mapping[str, Any])
30+
_ReadableBuffer = Union[bytes, memoryview, "mmap", "array"]

doc/examples/type_hints.rst

+2-2
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ type of document object returned when decoding BSON documents.
2020
Due to `limitations in mypy`_, the default
2121
values for generic document types are not yet provided (they will eventually be ``Dict[str, any]``).
2222

23-
For a larger set of examples that use types, see the PyMongo `test_mypy module`_.
23+
For a larger set of examples that use types, see the PyMongo `test_typing module`_.
2424

2525
If you would like to opt out of using the provided types, add the following to
2626
your `mypy config`_: ::
@@ -326,5 +326,5 @@ Another example is trying to set a value on a :class:`~bson.raw_bson.RawBSONDocu
326326
.. _mypy: https://mypy.readthedocs.io/en/stable/cheat_sheet_py3.html
327327
.. _limitations in mypy: https://github.com/python/mypy/issues/3737
328328
.. _mypy config: https://mypy.readthedocs.io/en/stable/config_file.html
329-
.. _test_mypy module: https://github.com/mongodb/mongo-python-driver/blob/master/test/test_mypy.py
329+
.. _test_typing module: https://github.com/mongodb/mongo-python-driver/blob/master/test/test_typing.py
330330
.. _schema validation: https://www.mongodb.com/docs/manual/core/schema-validation/#when-to-use-schema-validation

mypy.ini

+1-1
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ ignore_missing_imports = True
3232
[mypy-snappy.*]
3333
ignore_missing_imports = True
3434

35-
[mypy-test.test_mypy]
35+
[mypy-test.test_typing]
3636
warn_unused_ignores = True
3737

3838
[mypy-winkerberos.*]

pymongo/collection.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@
7272
InsertOneResult,
7373
UpdateResult,
7474
)
75-
from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
75+
from pymongo.typings import _CollationIn, _DocumentType, _DocumentTypeArg, _Pipeline
7676
from pymongo.write_concern import WriteConcern
7777

7878
_FIND_AND_MODIFY_DOC_FIELDS = {"value": 1}
@@ -103,6 +103,7 @@ class ReturnDocument(object):
103103

104104

105105
if TYPE_CHECKING:
106+
import bson
106107
from pymongo.client_session import ClientSession
107108
from pymongo.database import Database
108109
from pymongo.read_concern import ReadConcern
@@ -116,7 +117,7 @@ def __init__(
116117
database: "Database[_DocumentType]",
117118
name: str,
118119
create: Optional[bool] = False,
119-
codec_options: Optional[CodecOptions] = None,
120+
codec_options: Optional["CodecOptions[_DocumentTypeArg]"] = None,
120121
read_preference: Optional[_ServerMode] = None,
121122
write_concern: Optional[WriteConcern] = None,
122123
read_concern: Optional["ReadConcern"] = None,
@@ -394,7 +395,7 @@ def database(self) -> "Database[_DocumentType]":
394395

395396
def with_options(
396397
self,
397-
codec_options: Optional[CodecOptions] = None,
398+
codec_options: Optional["bson.CodecOptions[_DocumentTypeArg]"] = None,
398399
read_preference: Optional[_ServerMode] = None,
399400
write_concern: Optional[WriteConcern] = None,
400401
read_concern: Optional["ReadConcern"] = None,

pymongo/database.py

+9-8
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@
2929
cast,
3030
)
3131

32-
from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
32+
from bson.codec_options import DEFAULT_CODEC_OPTIONS
3333
from bson.dbref import DBRef
3434
from bson.son import SON
3535
from bson.timestamp import Timestamp
@@ -41,7 +41,7 @@
4141
from pymongo.common import _ecc_coll_name, _ecoc_coll_name, _esc_coll_name
4242
from pymongo.errors import CollectionInvalid, InvalidName
4343
from pymongo.read_preferences import ReadPreference, _ServerMode
44-
from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
44+
from pymongo.typings import _CollationIn, _DocumentType, _DocumentTypeArg, _Pipeline
4545

4646

4747
def _check_name(name):
@@ -55,6 +55,7 @@ def _check_name(name):
5555

5656

5757
if TYPE_CHECKING:
58+
import bson
5859
import bson.codec_options
5960
from pymongo.client_session import ClientSession
6061
from pymongo.mongo_client import MongoClient
@@ -72,7 +73,7 @@ def __init__(
7273
self,
7374
client: "MongoClient[_DocumentType]",
7475
name: str,
75-
codec_options: Optional[CodecOptions] = None,
76+
codec_options: Optional["bson.CodecOptions[_DocumentTypeArg]"] = None,
7677
read_preference: Optional[_ServerMode] = None,
7778
write_concern: Optional["WriteConcern"] = None,
7879
read_concern: Optional["ReadConcern"] = None,
@@ -152,7 +153,7 @@ def name(self) -> str:
152153

153154
def with_options(
154155
self,
155-
codec_options: Optional[CodecOptions] = None,
156+
codec_options: Optional["bson.CodecOptions[_DocumentTypeArg]"] = None,
156157
read_preference: Optional[_ServerMode] = None,
157158
write_concern: Optional["WriteConcern"] = None,
158159
read_concern: Optional["ReadConcern"] = None,
@@ -239,7 +240,7 @@ def __getitem__(self, name: str) -> "Collection[_DocumentType]":
239240
def get_collection(
240241
self,
241242
name: str,
242-
codec_options: Optional[CodecOptions] = None,
243+
codec_options: Optional["bson.CodecOptions[_DocumentTypeArg]"] = None,
243244
read_preference: Optional[_ServerMode] = None,
244245
write_concern: Optional["WriteConcern"] = None,
245246
read_concern: Optional["ReadConcern"] = None,
@@ -295,7 +296,7 @@ def get_collection(
295296
def create_collection(
296297
self,
297298
name: str,
298-
codec_options: Optional[CodecOptions] = None,
299+
codec_options: Optional["bson.CodecOptions[_DocumentTypeArg]"] = None,
299300
read_preference: Optional[_ServerMode] = None,
300301
write_concern: Optional["WriteConcern"] = None,
301302
read_concern: Optional["ReadConcern"] = None,
@@ -976,7 +977,7 @@ def _drop_helper(self, name, session=None, comment=None):
976977
@_csot.apply
977978
def drop_collection(
978979
self,
979-
name_or_collection: Union[str, Collection],
980+
name_or_collection: Union[str, Collection[_DocumentTypeArg]],
980981
session: Optional["ClientSession"] = None,
981982
comment: Optional[Any] = None,
982983
encrypted_fields: Optional[Mapping[str, Any]] = None,
@@ -1068,7 +1069,7 @@ def drop_collection(
10681069

10691070
def validate_collection(
10701071
self,
1071-
name_or_collection: Union[str, Collection],
1072+
name_or_collection: Union[str, Collection[_DocumentTypeArg]],
10721073
scandata: bool = False,
10731074
full: bool = False,
10741075
session: Optional["ClientSession"] = None,

pymongo/message.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import random
2525
import struct
2626
from io import BytesIO as _BytesIO
27-
from typing import Any, Dict, NoReturn
27+
from typing import Any, Mapping, NoReturn
2828

2929
import bson
3030
from bson import CodecOptions, _decode_selective, _dict_to_bson, _make_c_string, encode
@@ -81,7 +81,7 @@
8181
}
8282
_FIELD_MAP = {"insert": "documents", "update": "updates", "delete": "deletes"}
8383

84-
_UNICODE_REPLACE_CODEC_OPTIONS: "CodecOptions[Dict[str, Any]]" = CodecOptions(
84+
_UNICODE_REPLACE_CODEC_OPTIONS: "CodecOptions[Mapping[str, Any]]" = CodecOptions(
8585
unicode_decode_error_handler="replace"
8686
)
8787

0 commit comments

Comments
 (0)