PYTHON-3494 Improve Documentation Surrounding Type-Checking "_id" (#1104)

doc/examples/type_hints.rst

@@ -97,6 +97,7 @@ You can use :py:class:`~typing.TypedDict` (Python 3.8+) when using a well-define
 These methods automatically add an "_id" field.
 
 .. doctest::
+  :pyversion: >= 3.8
 
   >>> from typing import TypedDict
   >>> from pymongo import MongoClient
@@ -111,14 +112,73 @@ These methods automatically add an "_id" field.
   >>> result = collection.find_one({"name": "Jurassic Park"})
   >>> assert result is not None
   >>> assert result["year"] == 1993
-  >>> # This will not be type checked, despite being present, because it is added by PyMongo.
-  >>> assert type(result["_id"]) == ObjectId
+  >>> # This will raise a type-checking error, despite being present, because it is added by PyMongo.
+  >>> assert result["_id"]  # type:ignore[typeddict-item]
+
+Modeling Document Types with TypedDict
+--------------------------------------
+
+You can use :py:class:`~typing.TypedDict` (Python 3.8+) to model structured data.
+As noted above, PyMongo will automatically add an `_id` field if it is not present. This also applies to TypedDict.
+There are three approaches to this:
+
+  1. Do not specify `_id` at all. It will be inserted automatically, and can be retrieved at run-time, but will yield a type-checking error unless explicitly ignored.
+
+  2. Specify `_id` explicitly. This will mean that every instance of your custom TypedDict class will have to pass a value for `_id`.
+
+  3. Make use of :py:class:`~typing.NotRequired`. This has the flexibility of option 1, but with the ability to access the `_id` field without causing a type-checking error.
+
+Note: to use :py:class:`~typing.TypedDict` and :py:class:`~typing.NotRequired` in earlier versions of Python (<3.8, <3.11), use the `typing_extensions` package.
+
+.. doctest:: typed-dict-example
+  :pyversion: >= 3.11
+
+  >>> from typing import TypedDict, NotRequired
+  >>> from pymongo import MongoClient
+  >>> from pymongo.collection import Collection
+  >>> from bson import ObjectId
+  >>> class Movie(TypedDict):
+  ...       name: str
+  ...       year: int
+  ...
+  >>> class ExplicitMovie(TypedDict):
+  ...       _id: ObjectId
+  ...       name: str
+  ...       year: int
+  ...
+  >>> class NotRequiredMovie(TypedDict):
+  ...       _id: NotRequired[ObjectId]
+  ...       name: str
+  ...       year: int
+  ...
+  >>> client: MongoClient = MongoClient()
+  >>> collection: Collection[Movie] = client.test.test
+  >>> inserted = collection.insert_one(Movie(name="Jurassic Park", year=1993))
+  >>> result = collection.find_one({"name": "Jurassic Park"})
+  >>> assert result is not None
+  >>> # This will yield a type-checking error, despite being present, because it is added by PyMongo.
+  >>> assert result["_id"]  # type:ignore[typeddict-item]
+  >>> collection: Collection[ExplicitMovie] = client.test.test
+  >>> # Note that the _id keyword argument must be supplied
+  >>> inserted = collection.insert_one(ExplicitMovie(_id=ObjectId(), name="Jurassic Park", year=1993))
+  >>> result = collection.find_one({"name": "Jurassic Park"})
+  >>> assert result is not None
+  >>> # This will not raise a type-checking error.
+  >>> assert result["_id"]
+  >>> collection: Collection[NotRequiredMovie] = client.test.test
+  >>> # Note the lack of _id, similar to the first example
+  >>> inserted = collection.insert_one(NotRequiredMovie(name="Jurassic Park", year=1993))
+  >>> result = collection.find_one({"name": "Jurassic Park"})
+  >>> assert result is not None
+  >>> # This will not raise a type-checking error, despite not being provided explicitly.
+  >>> assert result["_id"]
+
 
 Typed Database
 --------------
 
 While less common, you could specify that the documents in an entire database
-match a well-defined shema using :py:class:`~typing.TypedDict` (Python 3.8+).
+match a well-defined schema using :py:class:`~typing.TypedDict` (Python 3.8+).
 
 
 .. doctest::

test/test_mypy.py

@@ -20,14 +20,30 @@ import unittest
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List
 
 try:
-    from typing_extensions import TypedDict
+    from typing_extensions import NotRequired, TypedDict
 
-    class Movie(TypedDict):  # type: ignore[misc]
+    from bson import ObjectId
+
+    class Movie(TypedDict):
         name: str
         year: int
 
-except ImportError:
-    TypedDict = None
+    class MovieWithId(TypedDict):
+        _id: ObjectId
+        name: str
+        year: int
+
+    class ImplicitMovie(TypedDict):
+        _id: NotRequired[ObjectId]
+        name: str
+        year: int
+
+except ImportError as exc:
+    Movie = dict  # type:ignore[misc,assignment]
+    ImplicitMovie = dict  # type: ignore[assignment,misc]
+    MovieWithId = dict  # type: ignore[assignment,misc]
+    TypedDict = None  # type: ignore[assignment]
+    NotRequired = None  # type: ignore[assignment]
 
 
 try:
@@ -324,6 +340,43 @@ class TestDocumentType(unittest.TestCase):
         )
         coll.insert_many([bad_movie])
 
+    @only_type_check
+    def test_typeddict_explicit_document_type(self) -> None:
+        out = MovieWithId(_id=ObjectId(), name="THX-1138", year=1971)
+        assert out is not None
+        # This should fail because the output is a Movie.
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert out["_id"]
+
+    # This should work the same as the test above, but this time using NotRequired to allow
+    # automatic insertion of the _id field by insert_one.
+    @only_type_check
+    def test_typeddict_not_required_document_type(self) -> None:
+        out = ImplicitMovie(name="THX-1138", year=1971)
+        assert out is not None
+        # This should fail because the output is a Movie.
+        assert out["foo"]  # type:ignore[typeddict-item]
+        assert out["_id"]
+
+    @only_type_check
+    def test_typeddict_empty_document_type(self) -> None:
+        out = Movie(name="THX-1138", year=1971)
+        assert out is not None
+        # This should fail because the output is a Movie.
+        assert out["foo"]  # type:ignore[typeddict-item]
+        # This should fail because _id is not included in our TypedDict definition.
+        assert out["_id"]  # type:ignore[typeddict-item]
+
+    def test_typeddict_find_notrequired(self):
+        if NotRequired is None or ImplicitMovie is None:
+            raise unittest.SkipTest("Python 3.11+ is required to use NotRequired.")
+        client: MongoClient[ImplicitMovie] = rs_or_single_client()
+        coll = client.test.test
+        coll.insert_one(ImplicitMovie(name="THX-1138", year=1971))
+        out = coll.find_one({})
+        assert out is not None
+        assert out["_id"]
+
     @only_type_check
     def test_raw_bson_document_type(self) -> None:
         client = MongoClient(document_class=RawBSONDocument)