Add document.WtfFormMixin with wtf related functions

Author: insspb

flask_mongoengine/documents.py

@@ -1,11 +1,21 @@
 """Extended version of :mod:`mongoengine.document`."""
+import logging
+from typing import List, Optional, Type, Union
+
 import mongoengine
 from flask import abort
 from mongoengine.errors import DoesNotExist
 from mongoengine.queryset import QuerySet
 
+from flask_mongoengine.decorators import wtf_required
 from flask_mongoengine.pagination import ListFieldPagination, Pagination
 
+try:
+    from flask_mongoengine.wtf.models import ModelForm
+except ImportError:  # pragma: no cover
+    ModelForm = None
+logger = logging.getLogger("flask_mongoengine")
+
 
 class BaseQuerySet(QuerySet):
     """Extends :class:`~mongoengine.queryset.QuerySet` class with handly methods."""
@@ -43,7 +53,7 @@ def first_or_404(self, _message_404=None):
         """
         return self.first() or self._abort_404(_message_404)
 
-    def paginate(self, page, per_page, **kwargs):
+    def paginate(self, page, per_page):
         """
         Paginate the QuerySet with a certain number of docs per page
         and return docs for a given page.
@@ -64,8 +74,89 @@ def paginate_field(self, field_name, doc_id, page, per_page, total=None):
         )
 
 
-class Document(mongoengine.Document):
-    """Abstract document with extra helpers in the queryset class"""
+class WtfFormMixin:
+    """Special mixin, for form generation functions."""
+
+    @classmethod
+    def _get_fields_names(
+        cls: Union["WtfFormMixin", mongoengine.document.BaseDocument],
+        only: Optional[List[str]],
+        exclude: Optional[List[str]],
+    ):
+        """
+        Filter fields names for further form generation.
+
+        :param only:
+            An optional iterable with the property names that should be included in
+            the form. Only these properties will have fields.
+            Fields are always appear in provided order, this allows user to change form
+            fields ordering, without changing database model.
+        :param exclude:
+            An optional iterable with the property names that should be excluded
+            from the form. All other properties will have fields.
+            Fields are appears in order, defined in model, excluding provided fields
+            names. For adjusting fields ordering, use :attr:`only`.
+        """
+        field_names = cls._fields_ordered
+
+        if only:
+            field_names = [field for field in only if field in field_names]
+        elif exclude:
+            field_names = [field for field in field_names if field not in exclude]
+
+        return field_names
+
+    @classmethod
+    @wtf_required
+    def to_wtf_form(
+        cls: Union["WtfFormMixin", mongoengine.document.BaseDocument],
+        base_class: Type[ModelForm] = ModelForm,
+        only: Optional[List[str]] = None,
+        exclude: Optional[List[str]] = None,
+        field_args=None,
+    ) -> Type[ModelForm]:
+        """
+        Generate WTForm from Document model.
+
+        :param base_class:
+            Base form class to extend from. Must be a :class:`.ModelForm` subclass.
+        :param only:
+            An optional iterable with the property names that should be included in
+            the form. Only these properties will have fields.
+            Fields are always appear in provided order, this allows user to change form
+            fields ordering, without changing database model.
+        :param exclude:
+            An optional iterable with the property names that should be excluded
+            from the form. All other properties will have fields.
+            Fields are appears in order, defined in model, excluding provided fields
+            names. For adjusting fields ordering, use :attr:`only`.
+        :param field_args:
+            An optional dictionary of field names mapping to keyword arguments used
+            to construct each field object.
+        """
+        form_fields_dict = {}
+        fields_names = cls._get_fields_names(only, exclude)
+
+        for field_name in fields_names:
+            # noinspection PyUnresolvedReferences
+            field_class = cls._fields[field_name]
+            try:
+                form_fields_dict[field_name] = field_class.to_wtf_field(
+                    field_args.get(field_name)
+                )
+            except AttributeError:
+                logger.warning(
+                    f"Field {field_name} ignored, field type does not have "
+                    f".to_wtf_field() method."
+                )
+
+        form_fields_dict["model_class"] = cls
+        # noinspection PyTypeChecker
+        return type(f"{cls.__name__}Form", (base_class,), form_fields_dict)
+
+
+class Document(WtfFormMixin, mongoengine.Document):
+    """Abstract Document with QuerySet and WTForms extra helpers."""
 
     meta = {"abstract": True, "queryset_class": BaseQuerySet}
 
@@ -79,7 +170,19 @@ def paginate_field(self, field_name, page, per_page, total=None):
         )
 
 
-class DynamicDocument(mongoengine.DynamicDocument):
-    """Abstract Dynamic document with extra helpers in the queryset class"""
+class DynamicDocument(WtfFormMixin, mongoengine.DynamicDocument):
+    """Abstract DynamicDocument with QuerySet and WTForms extra helpers."""
 
     meta = {"abstract": True, "queryset_class": BaseQuerySet}
+
+
+class EmbeddedDocument(WtfFormMixin, mongoengine.EmbeddedDocument):
+    """Abstract EmbeddedDocument document with extra WTForms helpers."""
+
+    meta = {"abstract": True}
+
+
+class DynamicEmbeddedDocument(WtfFormMixin, mongoengine.DynamicEmbeddedDocument):
+    """Abstract DynamicEmbeddedDocument document with extra WTForms helpers."""
+
+    meta = {"abstract": True}

tests/test_db_fields.py

@@ -5,7 +5,7 @@
 from mongoengine import fields as base_fields
 from pytest_mock import MockerFixture
 
-from flask_mongoengine import db_fields
+from flask_mongoengine import db_fields, documents
 
 
 @pytest.fixture
@@ -15,6 +15,77 @@ def local_app(app):
         yield app
 
 
+class TestWtfFormMixin:
+    @pytest.fixture(
+        params=[
+            documents.Document,
+            documents.DynamicDocument,
+            documents.DynamicEmbeddedDocument,
+            documents.EmbeddedDocument,
+        ]
+    )
+    def TempDocument(self, request):
+        """Cool fixture"""
+
+        class Model(request.param):
+            """
+            Temp document model for most of the tests.
+
+            field_one by design cannot be converted to FormField.
+            """
+
+            field_one = base_fields.StringField()
+            field_two = db_fields.StringField()
+            field_three = db_fields.StringField()
+            field_four = db_fields.StringField()
+            field_five = db_fields.StringField()
+
+        return Model
+
+    def test__get_fields_names__is_called_by_to_wtf_form_call(
+        self, TempDocument, mocker: MockerFixture
+    ):
+        get_fields_names_spy = mocker.patch.object(
+            documents.WtfFormMixin, "_get_fields_names", autospec=True
+        )
+        TempDocument.to_wtf_form()
+        get_fields_names_spy.assert_called_once()
+
+    def test__get_fields_names__hold_correct_fields_ordering_for_only(
+        self, TempDocument
+    ):
+
+        field_names = TempDocument._get_fields_names(
+            only=["field_five", "field_one"], exclude=None
+        )
+        assert field_names == ["field_five", "field_one"]
+
+    def test__get_fields_names__hold_correct_fields_ordering_for_exclude(
+        self, TempDocument
+    ):
+        field_names = TempDocument._get_fields_names(
+            only=None, exclude=["id", "field_five", "field_one"]
+        )
+        assert field_names == ["field_two", "field_three", "field_four"]
+
+    def test__to_wtf_form__is_called_by_mixin_child_model(
+        self, TempDocument, caplog, mocker: MockerFixture
+    ):
+        to_wtf_spy = mocker.patch.object(
+            documents.WtfFormMixin, "to_wtf_form", autospec=True
+        )
+        TempDocument.to_wtf_form()
+        to_wtf_spy.assert_called_once()
+
+    def test__to_wtf_form__logs_error(self, caplog, TempDocument):
+        TempDocument.to_wtf_form()
+
+        # Check error logging
+        assert (
+            "Field field_one ignored, field type does not have .to_wtf_field() method."
+        ) in caplog.messages
+
+
 class TestWtfFieldMixin:
     def test__init__set_additional_instance_arguments(self, db, mocker: MockerFixture):
         checker_spy = mocker.spy(db_fields.WtfFieldMixin, "_ensure_callable_or_list")