PYTHON-3090 Clean up Database Command Typing (#879)

blink1073 · web-flow · commit a61ea0660a0f · 2022-03-02T13:10:15.000-06:00
diff --git a/.github/workflows/test-python.yml b/.github/workflows/test-python.yml
@@ -64,4 +64,4 @@ jobs:
       - name: Run mypy
         run: |
           mypy --install-types --non-interactive bson gridfs tools pymongo
-          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 --exclude "test/mypy_fails/*.*" test
+          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
diff --git a/bson/codec_options.py b/bson/codec_options.py
@@ -115,6 +115,7 @@ class TypeCodec(TypeEncoder, TypeDecoder):
 
 _Codec = Union[TypeEncoder, TypeDecoder, TypeCodec]
 _Fallback = Callable[[Any], Any]
+_DocumentClass = Union[Type[MutableMapping], Type["RawBSONDocument"]]
 
 
 class TypeRegistry(object):
@@ -293,7 +294,7 @@ class CodecOptions(_options_base):
 
     def __new__(
         cls: Type["CodecOptions"],
-        document_class: Union[Type[MutableMapping], Type["RawBSONDocument"]] = dict,
+        document_class: _DocumentClass = dict,
         tz_aware: bool = False,
         uuid_representation: Optional[int] = UuidRepresentation.UNSPECIFIED,
         unicode_decode_error_handler: Optional[str] = "strict",
diff --git a/mypy.ini b/mypy.ini
@@ -32,9 +32,8 @@ ignore_missing_imports = True
 [mypy-snappy.*]
 ignore_missing_imports = True
 
-[mypy-test.*]
-allow_redefinition = true
-allow_untyped_globals = true
+[mypy-test.test_mypy]
+warn_unused_ignores = false
 
 [mypy-winkerberos.*]
 ignore_missing_imports = True
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -24,6 +24,7 @@
     Optional,
     Sequence,
     Union,
+    cast,
 )
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS, CodecOptions
@@ -37,7 +38,7 @@
 from pymongo.command_cursor import CommandCursor
 from pymongo.errors import CollectionInvalid, InvalidName
 from pymongo.read_preferences import ReadPreference, _ServerMode
-from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
+from pymongo.typings import _CollationIn, _DocumentOut, _DocumentType, _Pipeline
 
 
 def _check_name(name):
@@ -620,7 +621,7 @@ def command(
         session: Optional["ClientSession"] = None,
         comment: Optional[Any] = None,
         **kwargs: Any,
-    ) -> Dict[str, Any]:
+    ) -> _DocumentOut:
         """Issue a MongoDB command.
 
         Send command `command` to the database and return the
@@ -974,7 +975,7 @@ def validate_collection(
         if background is not None:
             cmd["background"] = background
 
-        result = self.command(cmd, session=session)
+        result = cast(dict, self.command(cmd, session=session))
 
         valid = True
         # Pre 1.9 results
diff --git a/pymongo/encryption.py b/pymongo/encryption.py
@@ -290,7 +290,7 @@ def _get_internal_client(encrypter, mongo_client):
         db, coll = opts._key_vault_namespace.split(".", 1)
         key_vault_coll = key_vault_client[db][coll]
 
-        mongocryptd_client = MongoClient(
+        mongocryptd_client: MongoClient = MongoClient(
             opts._mongocryptd_uri, connect=False, serverSelectionTimeoutMS=_MONGOCRYPTD_TIMEOUT_MS
         )
 
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -120,7 +120,7 @@ def __init__(
         self,
         host: Optional[Union[str, Sequence[str]]] = None,
         port: Optional[int] = None,
-        document_class: Type[_DocumentType] = dict,
+        document_class: Optional[Type[_DocumentType]] = None,
         tz_aware: Optional[bool] = None,
         connect: Optional[bool] = None,
         type_registry: Optional[TypeRegistry] = None,
@@ -652,7 +652,7 @@ def __init__(
         self.__init_kwargs: Dict[str, Any] = {
             "host": host,
             "port": port,
-            "document_class": document_class,
+            "document_class": document_class or dict,
             "tz_aware": tz_aware,
             "connect": connect,
             "type_registry": type_registry,
@@ -676,7 +676,7 @@ def __init__(
 
         # Parse options passed as kwargs.
         keyword_opts = common._CaseInsensitiveDictionary(kwargs)
-        keyword_opts["document_class"] = document_class
+        keyword_opts["document_class"] = document_class or dict
 
         seeds = set()
         username = None
@@ -1717,8 +1717,11 @@ def server_info(self, session: Optional[client_session.ClientSession] = None) ->
         .. versionchanged:: 3.6
            Added ``session`` parameter.
         """
-        return self.admin.command(
-            "buildinfo", read_preference=ReadPreference.PRIMARY, session=session
+        return cast(
+            dict,
+            self.admin.command(
+                "buildinfo", read_preference=ReadPreference.PRIMARY, session=session
+            ),
         )
 
     def list_databases(
diff --git a/pymongo/typings.py b/pymongo/typings.py
@@ -16,7 +16,6 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
     Mapping,
     MutableMapping,
     Optional,
@@ -36,6 +35,5 @@
 _CollationIn = Union[Mapping[str, Any], "Collation"]
 _DocumentIn = Union[MutableMapping[str, Any], "RawBSONDocument"]
 _Pipeline = Sequence[Mapping[str, Any]]
-_DocumentType = TypeVar(
-    "_DocumentType", Mapping[str, Any], MutableMapping[str, Any], Dict[str, Any]
-)
+_DocumentOut = _DocumentIn
+_DocumentType = TypeVar("_DocumentType", bound=Mapping[str, Any])
diff --git a/test/mypy_fails/raw_bson_document.py b/test/mypy_fails/raw_bson_document.py
@@ -0,0 +1,13 @@
+from bson.raw_bson import RawBSONDocument
+from pymongo import MongoClient
+
+client = MongoClient(document_class=RawBSONDocument)
+coll = client.test.test
+doc = {"my": "doc"}
+coll.insert_one(doc)
+retreived = coll.find_one({"_id": doc["_id"]})
+assert retreived is not None
+assert len(retreived.raw) > 0
+retreived[
+    "foo"
+] = "bar"  # error: Unsupported target for indexed assignment ("RawBSONDocument")  [index]
diff --git a/test/mypy_fails/typedict_client.py b/test/mypy_fails/typedict_client.py
@@ -0,0 +1,18 @@
+from typing import TypedDict
+
+from pymongo import MongoClient
+
+
+class Movie(TypedDict):
+    name: str
+    year: int
+
+
+client: MongoClient[Movie] = MongoClient()
+coll = client.test.test
+retreived = coll.find_one({"_id": "foo"})
+assert retreived is not None
+assert retreived["year"] == 1
+assert (
+    retreived["name"] == 2
+)  # error: Non-overlapping equality check (left operand type: "str", right operand type: "Literal[2]")  [comparison-overlap]
diff --git a/test/test_client.py b/test/test_client.py
@@ -709,7 +709,7 @@ def test_repr(self):
         # Used to test 'eval' below.
         import bson  # noqa: F401
 
-        client = MongoClient(
+        client = MongoClient(  # type: ignore[type-var]
             "mongodb://localhost:27017,localhost:27018/?replicaSet=replset"
             "&connectTimeoutMS=12345&w=1&wtimeoutms=100",
             connect=False,
diff --git a/test/test_mypy.py b/test/test_mypy.py
@@ -17,17 +17,32 @@
 
 import os
 import unittest
-from typing import Any, Dict, Iterable, List
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List
+
+try:
+    from typing import TypedDict  # type: ignore[attr-defined]
+
+    # Not available in Python 3.6 and Python 3.7
+    class Movie(TypedDict):  # type: ignore[misc]
+        name: str
+        year: int
+
+except ImportError:
+    TypeDict = None
+
 
 try:
     from mypy import api
 except ImportError:
-    api = None
+    api = None  # type: ignore[assignment]
 
 from test import IntegrationTest
+from test.utils import rs_or_single_client
 
+from bson.raw_bson import RawBSONDocument
 from bson.son import SON
 from pymongo.collection import Collection
+from pymongo.mongo_client import MongoClient
 from pymongo.operations import InsertOne
 
 TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "mypy_fails")
@@ -48,6 +63,8 @@ def ensure_mypy_fails(self, filename: str) -> None:
 
     def test_mypy_failures(self) -> None:
         for filename in get_tests():
+            if filename == "typeddict_client.py" and TypedDict is None:
+                continue
             with self.subTest(filename=filename):
                 self.ensure_mypy_fails(filename)
 
@@ -87,6 +104,66 @@ def test_bulk_write(self) -> None:
         result = self.coll.bulk_write(requests)
         self.assertTrue(result.acknowledged)
 
+    def test_command(self) -> None:
+        result = self.client.admin.command("ping")
+        items = result.items()
+
+    def test_list_collections(self) -> None:
+        cursor = self.client.test.list_collections()
+        value = cursor.next()
+        items = value.items()
+
+    def test_list_databases(self) -> None:
+        cursor = self.client.list_databases()
+        value = cursor.next()
+        value.items()
+
+    def test_default_document_type(self) -> None:
+        client = rs_or_single_client()
+        coll = client.test.test
+        doc = {"my": "doc"}
+        coll.insert_one(doc)
+        retreived = coll.find_one({"_id": doc["_id"]})
+        assert retreived is not None
+        retreived["a"] = 1
+
+    def test_explicit_document_type(self) -> None:
+        if not TYPE_CHECKING:
+            raise unittest.SkipTest("Do not use raw MongoClient")
+        client: MongoClient[Dict[str, Any]] = MongoClient()
+        coll = client.test.test
+        retreived = coll.find_one({"_id": "foo"})
+        assert retreived is not None
+        retreived["a"] = 1
+
+    def test_typeddict_document_type(self) -> None:
+        if not TYPE_CHECKING:
+            raise unittest.SkipTest("Do not use raw MongoClient")
+        client: MongoClient[Movie] = MongoClient()
+        coll = client.test.test
+        retreived = coll.find_one({"_id": "foo"})
+        assert retreived is not None
+        assert retreived["year"] == 1
+        assert retreived["name"] == "a"
+
+    def test_raw_bson_document_type(self) -> None:
+        if not TYPE_CHECKING:
+            raise unittest.SkipTest("Do not use raw MongoClient")
+        client = MongoClient(document_class=RawBSONDocument)
+        coll = client.test.test
+        retreived = coll.find_one({"_id": "foo"})
+        assert retreived is not None
+        assert len(retreived.raw) > 0
+
+    def test_son_document_type(self) -> None:
+        if not TYPE_CHECKING:
+            raise unittest.SkipTest("Do not use raw MongoClient")
+        client = MongoClient(document_class=SON[str, Any])
+        coll = client.test.test
+        retreived = coll.find_one({"_id": "foo"})
+        assert retreived is not None
+        retreived["a"] = 1
+
     def test_aggregate_pipeline(self) -> None:
         coll3 = self.client.test.test3
         coll3.insert_many(

Original file line number	Diff line number	Diff line change
`@@ -290,7 +290,7 @@ def _get_internal_client(encrypter, mongo_client):`
`290`	`290`	`db, coll = opts._key_vault_namespace.split(".", 1)`
`291`	`291`	`key_vault_coll = key_vault_client[db][coll]`
`292`	`292`
`293`		`- mongocryptd_client = MongoClient(`
	`293`	`+ mongocryptd_client: MongoClient = MongoClient(`
`294`	`294`	`opts._mongocryptd_uri, connect=False, serverSelectionTimeoutMS=_MONGOCRYPTD_TIMEOUT_MS`
`295`	`295`	`)`
`296`	`296`