From bb3c6f264094b9f5c2ce78a4eba45b8b3b11c049 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C4=90=E1=BB=97=20Ng=E1=BB=8Dc=20Ngo=E1=BA=A1n?=
 <ngocngoan060288@gmail.com>
Date: Tue, 27 Oct 2020 16:34:49 +0700
Subject: [PATCH] add django rest framework mongoengine support

---
 .../django_rest_framework_mongoengine.py      | 58 +++++++++++++++++++
 requirements/optionals.txt                    |  6 +-
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 drf_spectacular/contrib/django_rest_framework_mongoengine.py

diff --git a/drf_spectacular/contrib/django_rest_framework_mongoengine.py b/drf_spectacular/contrib/django_rest_framework_mongoengine.py
new file mode 100644
index 00000000..6c038ce4
--- /dev/null
+++ b/drf_spectacular/contrib/django_rest_framework_mongoengine.py
@@ -0,0 +1,58 @@
+from mongoengine import document, fields  # type: ignore
+from rest_framework.utils.model_meta import get_field_info
+from rest_framework_mongoengine import serializers, viewsets  # type: ignore
+
+from drf_spectacular.openapi import AutoSchema
+from drf_spectacular.plumbing import (
+    anyisinstance, build_basic_type, error, get_lib_doc_excludes, warn,
+)
+from drf_spectacular.types import OpenApiTypes
+
+
+class MongoEngineAutoSchema(AutoSchema):
+    def _map_model_field(self, model_field, direction):
+        assert isinstance(model_field, fields.BaseField)
+        # to get a fully initialized serializer field we use DRF's own init logic
+        try:
+            field_cls, field_kwargs = serializers.DocumentSerializer().build_field(
+                field_name=model_field.name,
+                info=get_field_info(model_field.model),
+                model_class=model_field.model,
+                nested_depth=0,
+            )
+            field = field_cls(**field_kwargs)
+        except:  # noqa
+            field = None
+
+        if field and not anyisinstance(field, [fields.EmbeddedDocumentField, fields.EmbeddedDocumentListField]):
+            return self._map_serializer_field(field, direction)
+        elif isinstance(model_field, fields.ReferenceField):
+            return self._map_model_field(model_field.target_field, direction)
+        elif hasattr(document, model_field.__class__.__name__):
+            # be graceful when the document field is not explicitly mapped to a serializer
+            internal_type = getattr(document, model_field.__class__.__name__)
+            field_cls = serializers.DocumentSerializer.serializer_field_mapping.get(internal_type)
+            if not field_cls:
+                warn(
+                    f'document field "{model_field.__class__.__name__}" has no mapping in '
+                    f'DocumentSerializer. it may be a deprecated field. defaulting to "string"'
+                )
+                return build_basic_type(OpenApiTypes.STR)
+            return self._map_serializer_field(field_cls(), direction)
+        else:
+            error(
+                f'could not resolve document field "{model_field}". failed to resolve through '
+                f'serializer_field_mapping, get_internal_type(), or any override mechanism. '
+                f'defaulting to "string"'
+            )
+            return build_basic_type(OpenApiTypes.STR)
+
+
+def get_mongoengine_extended_doc_excludes():
+    return get_lib_doc_excludes() + [
+        serializers.DocumentSerializer,
+        serializers.EmbeddedDocumentSerializer,
+        viewsets.ModelViewSet,
+        viewsets.GenericViewSet,
+        viewsets.ReadOnlyModelViewSet
+    ]
diff --git a/requirements/optionals.txt b/requirements/optionals.txt
index 7272a266..5c62eaad 100644
--- a/requirements/optionals.txt
+++ b/requirements/optionals.txt
@@ -6,4 +6,8 @@ django-polymorphic>=2.1
 django-rest-polymorphic>=0.1.8
 django-oauth-toolkit>=1.2.0
 djangorestframework-camel-case>=1.1.2
-django-filter>=2.3.0
\ No newline at end of file
+django-filter>=2.3.0
+pymongo>=3.11.0
+mongoengine>=0.20.0
+django-mongoengine>=0.4.2
+django-rest-framework-mongoengine>=3.4.1
\ No newline at end of file