From 5c843e60f6cd26ebc232dac66b56c626a5808e73 Mon Sep 17 00:00:00 2001
From: Serhii Tereshchenko <serg.partizan@gmail.com>
Date: Fri, 3 Nov 2023 10:41:07 +0200
Subject: [PATCH] fix(typing): Add static types for dynamic classes

Refs https://github.com/microsoft/pyright/issues/6309
---
 django_mongoengine/document.py | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

diff --git a/django_mongoengine/document.py b/django_mongoengine/document.py
index 58c5833..eca4056 100644
--- a/django_mongoengine/document.py
+++ b/django_mongoengine/document.py
@@ -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
@@ -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()
@@ -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):
+        ...