Awesome
mongo-types
Type stubs for mongoengine
.
Allows for autocomplete and static typing.
install
pip install mongo-types
Monkey patch mongoengine's QuerySet
so we can type it with a generic
argument at runtime:
import types
from mongoengine.queryset.queryset import QuerySet
def no_op(self, x):
return self
QuerySet.__class_getitem__ = types.MethodType(no_op, QuerySet)
usage
After installing and monkey patching, the types should work for the most part, but you'll probably need to change how you write some things.
getting objects
to work
By default, the base document is typed to not have an objects
property so
that each document can type it properly.
Here's a helper class that's useful for simple cases which don't modify the
QuerySet
.
from typing import Generic, Type, TypeVar
from mongoengine import QuerySet, Document
U = TypeVar("U", bound=Document)
class QuerySetManager(Generic[U]):
def __get__(self, instance: object, cls: Type[U]) -> QuerySet[U]:
return QuerySet(cls, cls._get_collection())
class Page(Document):
meta = {
"collection": "pages",
}
objects = QuerySetManager["Page"]()
organization = fields.StringField()
replacing usages of queryset_class
before:
from typing import Type
from mongoengine import QuerySet, Document
class PostQuerySet(QuerySet):
def for_org(self, *, org: str) -> QuerySet:
return self.filter(organization=org)
def exists(self) -> bool:
return self.count() > 0
class Post(Document):
meta = {
"collection": "posts",
"queryset_class": SMSLogQuerySet,
}
organization = fields.StringField()
# --snip--
after:
from typing import Type
from mongoengine import QuerySet, Document
class PostQuerySet(QuerySet["Post"]):
def for_org(self, *, org: str) -> QuerySet["Post"]:
return self.filter(organization=org)
def exists(self) -> bool:
return self.count() > 0
class QuerySetManager:
def __get__(self, instance: object, cls: Type[Post]) -> PostQuerySet:
return PostQuerySet(cls, cls._get_collection())
class Post(Document):
meta = {
"collection": "posts",
}
objects = QuerySetManager()
organization = fields.StringField()
# --snip--
replicating @queryset_manager
behavior
before:
from mongoengine import Document, QuerySet, queryset_manager, fields
class UserQuerySet(QuerySet):
def for_org(self, *, org: str) -> QuerySet:
return self.filter(organization=org)
class User(Document):
meta = {
"collection": "users",
"queryset_class": UserQuerySet,
}
is_active = fields.BooleanField()
# --snip--
@queryset_manager
def objects(self, queryset: QuerySet) -> QuerySet:
return queryset.filter(is_active=True)
@queryset_manager
def all_objects(self, queryset: QuerySet) -> QuerySet:
return queryset
maybe_user = User.all_objects.first()
after:
from __future__ import annotations
from typing import Type
from mongoengine import QuerySet, Document
class UserQuerySet(QuerySet["User"]):
def for_org(self, *, org: str) -> UserQuerySet:
return self.filter(organization=org)
class QuerySetManager:
def __get__(self, instance: object, cls: Type[User]) -> UserQuerySet:
return UserQuerySet(cls, cls._get_collection()).filter(is_active=True)
class User(Document):
meta = {
"collection": "users",
}
is_active = fields.BooleanField()
# --snip--
objects = QuerySetManager()
@classmethod
def all_objects(cls) -> UserQuerySet:
return UserQuerySet(cls, cls._get_collection())
maybe_user = User.all_objects().first()
fixing "Model" has no attribute "id"
Mongoengine will define an id
field for you automatically.
Mongo-types require you specify your id
explicitly so that
the types can be more strict.
class User(Document):
meta = {
"collection": "users",
}
# becomes
class User(Document):
meta = {
"collection": "users",
}
id = fields.StringField(db_field="_id", primary_key=True, default=default_id)
# or if you prefer ObjectIds
class User(Document):
meta = {
"collection": "users",
}
id = fields.ObjectIdField(db_field="_id", primary_key=True, default=ObjectId)
dev
poetry install
# run formatting, linting, and typechecking
s/lint
# build
poetry build -f wheel
# build and publish
poetry publish --build