Skip to content

Commit

Permalink
fix(typing): Add static types for dynamic classes
Browse files Browse the repository at this point in the history
  • Loading branch information
last-partizan committed Nov 3, 2023
1 parent 505d126 commit 5c843e6
Showing 1 changed file with 25 additions and 2 deletions.
27 changes: 25 additions & 2 deletions django_mongoengine/document.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,13 @@
# ruff: noqa: F811
from __future__ import annotations

from functools import partial
from typing import TYPE_CHECKING, Any

from bson.objectid import ObjectId
from django.db.models import Model
from django.db.models.base import ModelState
from mongoengine import DoesNotExist
from mongoengine import document as me
from mongoengine.base import metaclasses as mtc
from mongoengine.errors import FieldDoesNotExist
Expand Down Expand Up @@ -38,9 +43,12 @@ def __new__(cls, name, bases, attrs):


class DjangoFlavor:
objects = QuerySetManager()
_default_manager = QuerySetManager()
id: Any
objects: Any = QuerySetManager()
_meta: DocumentMetaWrapper
_default_manager: Any = QuerySetManager()
_get_pk_val = Model.__dict__["_get_pk_val"]
DoesNotExist: type[DoesNotExist]

def __init__(self, *args, **kwargs):
self._state = ModelState()
Expand Down Expand Up @@ -96,3 +104,18 @@ class DynamicEmbeddedDocument(
django_meta(mtc.DocumentMetaclass, DjangoFlavor, me.DynamicEmbeddedDocument)
):
swap_base = True


if TYPE_CHECKING:

class Document(DjangoFlavor, me.Document):
...

class DynamicDocument(DjangoFlavor, me.DynamicDocument):
...

class EmbeddedDocument(DjangoFlavor, me.EmbeddedDocument):
...

class DynamicEmbeddedDocument(DjangoFlavor, me.DynamicEmbeddedDocument):
...

0 comments on commit 5c843e6

Please sign in to comment.