Merge pull request #489 from MongoEngine/orm_module_simplification

Updated: ORM module functions split for better maintenance

Author: insspb

flask_mongoengine/__init__.py

@@ -5,16 +5,12 @@
 from mongoengine.errors import DoesNotExist
 from mongoengine.queryset import QuerySet
 
+from flask_mongoengine import db_fields
 from flask_mongoengine.connection import *
 from flask_mongoengine.json import override_json_encoder
 from flask_mongoengine.pagination import *
 from flask_mongoengine.sessions import *
 
-try:
-    from flask_mongoengine.wtf import db_fields
-except ImportError:
-    from mongoengine import fields as db_fields
-
 
 def current_mongoengine_instance():
     """Return a MongoEngine instance associated with current Flask app."""

flask_mongoengine/wtf/db_fields.py renamed to flask_mongoengine/db_fields.py

@@ -1,6 +1,4 @@
 """Responsible for mongoengine fields extension, if WTFForms integration used."""
-from mongoengine import fields
-
 __all__ = [
     "WtfFieldMixin",
     "BinaryField",
@@ -43,6 +41,16 @@
     "URLField",
     "UUIDField",
 ]
+from typing import Callable, List, Optional, Union
+
+from mongoengine import fields
+
+try:
+    from wtforms import fields as wtf_fields
+    from wtforms import validators as wtf_validators
+except ImportError:  # pragma: no cover
+    wtf_fields = None
+    wtf_validators = None
 
 
 class WtfFieldMixin:
@@ -58,28 +66,35 @@ class WtfFieldMixin:
     :param kwargs: keyword arguments silently bypassed to normal mongoengine fields
     """
 
-    def __init__(self, *, validators=None, filters=None, **kwargs):
+    def __init__(
+        self,
+        *,
+        validators: Optional[Union[List, Callable]] = None,
+        filters: Optional[Union[List, Callable]] = None,
+        **kwargs,
+    ):
         self.validators = self._ensure_callable_or_list(validators, "validators")
         self.filters = self._ensure_callable_or_list(filters, "filters")
 
         super().__init__(**kwargs)
 
-    def _ensure_callable_or_list(self, field, msg_flag):
+    @staticmethod
+    def _ensure_callable_or_list(argument, msg_flag: str) -> Optional[List]:
         """
-        Ensure the value submitted via field is either
-        a callable object to convert to list or it is
-        in fact a valid list value.
+        Ensure submitted argument value is a callable object or valid list value.
 
+        :param argument: Argument input to make verification on.
+        :param msg_flag: Argument string name for error message.
         """
-        if field is not None:
-            if callable(field):
-                field = [field]
-            else:
-                msg = "Argument '%s' must be a list value" % msg_flag
-                if not isinstance(field, list):
-                    raise TypeError(msg)
-
-        return field
+        if argument is None:
+            return None
+
+        if callable(argument):
+            return [argument]
+        elif not isinstance(argument, list):
+            raise TypeError(f"Argument '{msg_flag}' must be a list value")
+
+        return argument
 
 
 class BinaryField(WtfFieldMixin, fields.BinaryField):

flask_mongoengine/decorators.py

@@ -0,0 +1,25 @@
+import functools
+import logging
+
+try:
+    import wtforms  # noqa
+
+    wtf_installed = True
+except ImportError:  # pragma: no cover
+    wtf_installed = False
+
+logger = logging.getLogger("flask_mongoengine")
+
+
+def wtf_required(func):
+    """Special decorator to warn user on incorrect installation."""
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        if not wtf_installed:
+            logger.error(f"WTForms not installed. Function '{func.__name__}' aborted.")
+            return None
+
+        return func(*args, **kwargs)
+
+    return wrapped

flask_mongoengine/wtf/__init__.py

@@ -1,2 +1,2 @@
 """WTFForms integration module init file."""
-from flask_mongoengine.wtf.db_fields import *  # noqa
+from flask_mongoengine.wtf.orm import model_fields, model_form  # noqa

flask_mongoengine/wtf/models.py

@@ -1,10 +1,15 @@
+from typing import Type, Union
+
+import mongoengine
 from flask_wtf import FlaskForm
 from flask_wtf.form import _Auto
 
 
 class ModelForm(FlaskForm):
     """A WTForms mongoengine model form"""
 
+    model_class: Type[Union[mongoengine.Document, mongoengine.DynamicDocument]]
+
     def __init__(self, formdata=_Auto, **kwargs):
         self.instance = kwargs.pop("instance", None) or kwargs.get("obj")
         if self.instance and not formdata:

flask_mongoengine/wtf/orm.py

@@ -3,9 +3,11 @@
 """
 import decimal
 from collections import OrderedDict
+from typing import List, Optional, Type
 
 from bson import ObjectId
 from mongoengine import ReferenceField
+from mongoengine.base import BaseDocument, DocumentMetaclass
 from wtforms import fields as f
 from wtforms import validators
 
@@ -45,8 +47,8 @@ def __init__(self, converters=None):
 
         self.converters = converters
 
-    def convert(self, model, field, field_args):
-        kwargs = {
+    def _generate_convert_base_kwargs(self, field, field_args) -> dict:
+        kwargs: dict = {
             "label": getattr(field, "verbose_name", field.name),
             "description": getattr(field, "help_text", None) or "",
             "validators": getattr(field, "validators", None) or [],
@@ -56,42 +58,49 @@ def convert(self, model, field, field_args):
         if field_args:
             kwargs.update(field_args)
 
-        if kwargs["validators"]:
-            # Create a copy of the list since we will be modifying it, and if
-            # validators set as shared list between fields - duplicates/conflicts may
-            # be created.
-            kwargs["validators"] = list(kwargs["validators"])
-
+        # Create a copy of the lists since we will be modifying it, and if
+        # validators set as shared list between fields - duplicates/conflicts may
+        # be created.
+        kwargs["validators"] = list(kwargs["validators"])
+        kwargs["filters"] = list(kwargs["filters"])
         if field.required:
             kwargs["validators"].append(validators.InputRequired())
         else:
             kwargs["validators"].append(validators.Optional())
 
-        ftype = type(field).__name__
+        return kwargs
 
-        if field.choices:
-            kwargs["choices"] = field.choices
-
-            if ftype in self.converters:
-                kwargs["coerce"] = self.coerce(ftype)
-            multiple_field = kwargs.pop("multiple", False)
-            radio_field = kwargs.pop("radio", False)
-            if multiple_field:
-                return f.SelectMultipleField(**kwargs)
-            if radio_field:
-                return f.RadioField(**kwargs)
-            return f.SelectField(**kwargs)
+    def _process_convert_for_choice_fields(self, field, field_class, kwargs):
+        kwargs["choices"] = field.choices
+        kwargs["coerce"] = self.coerce(field_class)
+        if kwargs.pop("multiple", False):
+            return f.SelectMultipleField(**kwargs)
+        if kwargs.pop("radio", False):
+            return f.RadioField(**kwargs)
+        return f.SelectField(**kwargs)
 
+    def convert(self, model, field, field_args):
         if hasattr(field, "to_form_field"):
-            return field.to_form_field(model, kwargs)
+            return field.to_form_field(model, field_args)
+
+        field_class = type(field).__name__
+
+        if field_class not in self.converters:
+            raise NotImplementedError(
+                f"No converter for: {field_class}, exclude it from form generation."
+            )
+
+        kwargs = self._generate_convert_base_kwargs(field, field_args)
+
+        if field.choices:
+            return self._process_convert_for_choice_fields(field, field_class, kwargs)
 
         if hasattr(field, "field") and isinstance(field.field, ReferenceField):
             kwargs["label_modifier"] = getattr(
                 model, f"{field.name}_label_modifier", None
             )
 
-        if ftype in self.converters:
-            return self.converters[ftype](model, field, kwargs)
+        return self.converters[field_class](model, field, kwargs)
 
     @classmethod
     def _string_common(cls, model, field, kwargs):
@@ -237,46 +246,68 @@ def coerce(self, field_type):
         return coercions.get(field_type, str)
 
 
-def model_fields(model, only=None, exclude=None, field_args=None, converter=None):
+def _get_fields_names(
+    model,
+    only: Optional[List[str]],
+    exclude: Optional[List[str]],
+) -> List[str]:
     """
-    Generate a dictionary of fields for a given database model.
+    Filter fields names for further form generation.
 
-    See `model_form` docstring for description of parameters.
+    :param model: Source model class for fields list retrieval
+    :param only: If provided, only these field names will have fields definition.
+    :param exclude: If provided, field names will be excluded from fields definition.
+      All other field names will have fields.
     """
-    from mongoengine.base import BaseDocument, DocumentMetaclass
+    field_names = model._fields_ordered
 
-    if not isinstance(model, (BaseDocument, DocumentMetaclass)):
+    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 set(exclude)]
+
+    return field_names
+
+
+def model_fields(
+    model: Type[BaseDocument],
+    only: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    field_args=None,
+    converter=None,
+) -> OrderedDict:
+    """
+    Generate a dictionary of fields for a given database model.
+
+    See :func:`model_form` docstring for description of parameters.
+    """
+    if not issubclass(model, (BaseDocument, DocumentMetaclass)):
         raise TypeError("model must be a mongoengine Document schema")
 
     converter = converter or ModelConverter()
     field_args = field_args or {}
+    form_fields_dict = OrderedDict()
+    # noinspection PyTypeChecker
+    fields_names = _get_fields_names(model, only, exclude)
 
-    names = ((k, v.creation_counter) for k, v in model._fields.items())
-    field_names = [n[0] for n in sorted(names, key=lambda n: n[1])]
+    for field_name in fields_names:
+        # noinspection PyUnresolvedReferences
+        model_field = model._fields[field_name]
+        form_field = converter.convert(model, model_field, field_args.get(field_name))
+        if form_field is not None:
+            form_fields_dict[field_name] = form_field
 
-    if only:
-        field_names = [x for x in only if x in set(field_names)]
-    elif exclude:
-        field_names = [x for x in field_names if x not in set(exclude)]
-
-    field_dict = OrderedDict()
-    for name in field_names:
-        model_field = model._fields[name]
-        field = converter.convert(model, model_field, field_args.get(name))
-        if field is not None:
-            field_dict[name] = field
-
-    return field_dict
+    return form_fields_dict
 
 
 def model_form(
-    model,
-    base_class=ModelForm,
-    only=None,
-    exclude=None,
+    model: Type[BaseDocument],
+    base_class: Type[ModelForm] = ModelForm,
+    only: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
     field_args=None,
     converter=None,
-):
+) -> Type[ModelForm]:
     """
     Create a wtforms Form for a given mongoengine Document schema::
 
@@ -287,7 +318,7 @@ def model_form(
     :param model:
         A mongoengine Document schema class
     :param base_class:
-        Base form class to extend from. Must be a ``wtforms.Form`` subclass.
+        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.
@@ -299,8 +330,9 @@ def model_form(
         to construct each field object.
     :param converter:
         A converter to generate the fields based on the model properties. If
-        not set, ``ModelConverter`` is used.
+        not set, :class:`.ModelConverter` is used.
     """
     field_dict = model_fields(model, only, exclude, field_args, converter)
     field_dict["model_class"] = model
+    # noinspection PyTypeChecker
     return type(f"{model.__name__}Form", (base_class,), field_dict)