MongoEngine support

[gmacon]

pyrefly currently has special support for Pydantic and Django models, would folks be interested in support for MongoEngine? MongoEngine makes heavy use of metaclasses, so the normal inference is missing most of the useful information.

[rchen152]

Sorry for the slow response - this slipped through the cracks. Currently, we're not actively expanding library support beyond Pydantic and Django, but we'd be happy to consider targeted changes to support specific common patterns in MongoEngine - we've added a few of those on a case-by-case basis for other libraries. Could you go into a bit more detail about your use case and what type information is missing? That would help me figure out the scope and feasibility.

[gmacon]

Well, we use MongoEngine to model our database data types. A model looks like this:

import datetime

from mongoengine import DateTimeField, Document, StringField

class User(Document):
    email = StringField(required=True)
    first_name = StringField(max_length=50)
    last_name = StringField(max_length=50)
    created = DateTimeField(required=True, default=partial(datetime.datetime.now, datetime.utc))

methods coming directly from Document (like save) work fine as-is, but attributes added by the metaclass, like the objects attribute that gives you access to the QuerySet interface are missing. The implicit id field is also missing. Quickly sketching, I think the type stubs would be something like:

import datetime
from typing import ClassVar, Self

from bson import ObjectId
from mongoengine import Document, QuerySet

# Inherit from Document to get most of the methods
class User(Document):
    def __init__(
        self,
        *,
        id: ObjectId = ...,
        email: str,
        first_name: str | None = None,
        last_name: str | None = None,
        created: datetime.datetime = ...,
    ) -> None: ...

    # real fields
    id: ObjectId
    email: str
    first_name: str | None
    last_name: str | None
    created: datetime.datetime

    # special attributes added by metaclass
    objects: ClassVar[QuerySet[Self]]
    pk: ObjectId

The signature of __init__ is fiction. The real implementation just has

def __init__(self, *, **values) -> None: ...

in the base class. It even permits you to omit required fields and doesn't complain until you try to save the document. Though I think, from a typing perspective, it would be better to require required fields to be passed in the constructor, that's not strictly required.

QuerySet is a real class in MongoEngine, but it isn't actually generic. That's a fiction that's needed to correctly infer the type of user in a statement like user = User.objects.get(email="foo@example.com") as being User instead of just Document. (It's technically possible to customize the object produced by the objects property, but I've never seen anyone do that, and I don't think pyrefly needs to support that.)

There's also a translation needed from the various subclasses of Field into the attribute types, but I think adding annotations like

class User(Document):
    email: str = StringField(required=True)
    first_name: str | None = StringField()
    last_name: str | None = StringField()
    created: datetime = DateTimeField(required=True, default=partial(datetime.datetime.now, datetime.utc))

would be better because it provides a clean way to handle custom field types (of which I have a couple).

The pk attribute is a property accessor for the primary key. If you have a custom primary key, like:

class User(Document):
    email: str = StringField(required=True, primary_key=True, db_field="_id")

then this would give annotations like

class User(Document):
    email: str
    pk: str

   # ... other things as before

with no implicit id field and pk being the same type as the primary key.

There's also EmbeddedDocument and the associated EmbeddedDocumentField for deeply-structured objects, but I think those could likely be handled along with Document in a mostly generic way due to the common BaseDocument base class.

I haven't looked at Django in a really long time (we're a Flask shop), but I believe most of this is similar to the way the Django ORM works.

[rchen152]

Thank you for the details! From a typing perspective, this is indeed very similar to what we have for Django, so we may be able to extend that code to cover MongoEngine. I've opened #3335.