diff --git a/doc/examples/type_hints.rst b/doc/examples/type_hints.rst index cd178038a..38349038b 100644 --- a/doc/examples/type_hints.rst +++ b/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:: diff --git a/test/test_mypy.py b/test/test_mypy.py index a1e94937b..807f0e8ef 100644 --- a/test/test_mypy.py +++ b/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)