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 0be278d commit 09260f0
Showing 1 changed file with 23 additions and 3 deletions.
26 changes: 23 additions & 3 deletions django_mongoengine/document.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
# ruff: noqa: F811
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 document as me
from mongoengine import DoesNotExist, document as me
from mongoengine.base import metaclasses as mtc
from mongoengine.errors import FieldDoesNotExist

Expand Down Expand Up @@ -38,9 +40,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 +101,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 09260f0

Please sign in to comment.