Fix to_dict for IntEnum

cuenca-mx · Oct 7, 2023 · 0aaf4db · 0aaf4db
1 parent f6128be
commit 0aaf4db
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/mongoengine_plus/models/helpers.py b/mongoengine_plus/models/helpers.py
@@ -70,7 +70,7 @@ def mongo_to_dict(obj, exclude_fields: list = None) -> dict:
         elif isinstance(obj._fields[field_name], DictField):
             return_data[field_name] = data
         elif isinstance(obj._fields[field_name], EnumField):
-            return_data[field_name] = data.value if data else None
+            return_data[field_name] = data.value if data is not None else None
         elif isinstance(obj._fields[field_name], LazyReferenceField):
             return_data[f'{field_name}_uri'] = (
                 f'/{data._DBRef__collection}/{data.id}' if data else None

diff --git a/mongoengine_plus/version.py b/mongoengine_plus/version.py
@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = '0.1.1.dev0'
diff --git a/tests/models/test_helpers.py b/tests/models/test_helpers.py
@@ -30,6 +30,12 @@ class EnumType(Enum):
     member = 'name'
 
 
+class IntEnum(int, Enum):
+    exito = 0
+    invalid = 1
+
+
+
 class Embedded(EmbeddedDocument):
     meta = {'allow_inheritance': True}
     name = StringField()
@@ -59,6 +65,7 @@ class TestModel(Document):
     lazzy_field = LazyReferenceField(Reference)
     lazzy_list_field = ListField(LazyReferenceField(Reference))
     generic_lazzy_field = GenericLazyReferenceField()
+    code = EnumField(IntEnum)
 
     __test__ = False
 
@@ -71,6 +78,7 @@ def test_mongo_to_dict():
         lazzy_list_field=[reference],
         embedded_field=Embedded(name='Peter'),
         heritage_field=HeritageEmbedded(name='some', lastname='other'),
+        code=IntEnum.exito,
     )
     model.save()
     model_dict = mongo_to_dict(model, exclude_fields=['str_field'])
@@ -95,3 +103,4 @@ def test_mongo_to_dict():
         'name': 'some',
         'lastname': 'other',
     }
+    assert model_dict['code'] == 0