Merge commit master into bugfix-save-sharding

mongoengine/base/document.py

@@ -538,6 +538,9 @@ class BaseDocument:
         """Using _get_changed_fields iterate and remove any fields that
         are marked as changed.
         """
+        ReferenceField = _import_class("ReferenceField")
+        GenericReferenceField = _import_class("GenericReferenceField")
+
         for changed in self._get_changed_fields():
             parts = changed.split(".")
             data = self
@@ -550,7 +553,8 @@ class BaseDocument:
                 elif isinstance(data, dict):
                     data = data.get(part, None)
                 else:
-                    data = getattr(data, part, None)
+                    field_name = data._reverse_db_field_map.get(part, part)
+                    data = getattr(data, field_name, None)
 
                 if not isinstance(data, LazyReference) and hasattr(
                     data, "_changed_fields"
@@ -559,10 +563,40 @@ class BaseDocument:
                         continue
 
                     data._changed_fields = []
+                elif isinstance(data, (list, tuple, dict)):
+                    if hasattr(data, "field") and isinstance(
+                        data.field, (ReferenceField, GenericReferenceField)
+                    ):
+                        continue
+                    BaseDocument._nestable_types_clear_changed_fields(data)
 
         self._changed_fields = []
 
-    def _nestable_types_changed_fields(self, changed_fields, base_key, data):
+    @staticmethod
+    def _nestable_types_clear_changed_fields(data):
+        """Inspect nested data for changed fields
+
+        :param data: data to inspect for changes
+        """
+        Document = _import_class("Document")
+
+        # Loop list / dict fields as they contain documents
+        # Determine the iterator to use
+        if not hasattr(data, "items"):
+            iterator = enumerate(data)
+        else:
+            iterator = data.items()
+
+        for index_or_key, value in iterator:
+            if hasattr(value, "_get_changed_fields") and not isinstance(
+                value, Document
+            ):  # don't follow references
+                value._clear_changed_fields()
+            elif isinstance(value, (list, tuple, dict)):
+                BaseDocument._nestable_types_clear_changed_fields(value)
+
+    @staticmethod
+    def _nestable_types_changed_fields(changed_fields, base_key, data):
         """Inspect nested data for changed fields
 
         :param changed_fields: Previously collected changed fields
@@ -587,7 +621,9 @@ class BaseDocument:
                 changed = value._get_changed_fields()
                 changed_fields += ["{}{}".format(item_key, k) for k in changed if k]
             elif isinstance(value, (list, tuple, dict)):
-                self._nestable_types_changed_fields(changed_fields, item_key, value)
+                BaseDocument._nestable_types_changed_fields(
+                    changed_fields, item_key, value
+                )
 
     def _get_changed_fields(self):
         """Return a list of all fields that have explicitly been changed.

mongoengine/fields.py

@@ -87,6 +87,7 @@ __all__ = (
     "PolygonField",
     "SequenceField",
     "UUIDField",
+    "EnumField",
     "MultiPointField",
     "MultiLineStringField",
     "MultiPolygonField",
@@ -847,8 +848,7 @@ class DynamicField(BaseField):
     Used by :class:`~mongoengine.DynamicDocument` to handle dynamic data"""
 
     def to_mongo(self, value, use_db_field=True, fields=None):
-        """Convert a Python type to a MongoDB compatible type.
-        """
+        """Convert a Python type to a MongoDB compatible type."""
 
         if isinstance(value, str):
             return value
@@ -1622,6 +1622,70 @@ class BinaryField(BaseField):
         return super().prepare_query_value(op, self.to_mongo(value))
 
 
+class EnumField(BaseField):
+    """Enumeration Field. Values are stored underneath as is,
+    so it will only work with simple types (str, int, etc) that
+    are bson encodable
+     Example usage:
+    .. code-block:: python
+
+        class Status(Enum):
+            NEW = 'new'
+            DONE = 'done'
+
+        class ModelWithEnum(Document):
+            status = EnumField(Status, default=Status.NEW)
+
+        ModelWithEnum(status='done')
+        ModelWithEnum(status=Status.DONE)
+
+    Enum fields can be searched using enum or its value:
+    .. code-block:: python
+
+        ModelWithEnum.objects(status='new').count()
+        ModelWithEnum.objects(status=Status.NEW).count()
+
+    Note that choices cannot be set explicitly, they are derived
+    from the provided enum class.
+    """
+
+    def __init__(self, enum, **kwargs):
+        self._enum_cls = enum
+        if "choices" in kwargs:
+            raise ValueError(
+                "'choices' can't be set on EnumField, "
+                "it is implicitly set as the enum class"
+            )
+        kwargs["choices"] = list(self._enum_cls)
+        super().__init__(**kwargs)
+
+    def __set__(self, instance, value):
+        is_legal_value = value is None or isinstance(value, self._enum_cls)
+        if not is_legal_value:
+            try:
+                value = self._enum_cls(value)
+            except Exception:
+                pass
+        return super().__set__(instance, value)
+
+    def to_mongo(self, value):
+        if isinstance(value, self._enum_cls):
+            return value.value
+        return value
+
+    def validate(self, value):
+        if value and not isinstance(value, self._enum_cls):
+            try:
+                self._enum_cls(value)
+            except Exception as e:
+                self.error(str(e))
+
+    def prepare_query_value(self, op, value):
+        if value is None:
+            return value
+        return super().prepare_query_value(op, self.to_mongo(value))
+
+
 class GridFSError(Exception):
     pass