add enum field (#626)

* add enum field

* add decorator for asyncio

* fix enum typing, additional tests, add docs

* add more tests

Co-authored-by: collerek <collerek@gmail.com>

tests/test_model_definition/test_columns.py

@@ -1,4 +1,5 @@
 import datetime
+from enum import Enum
 
 import databases
 import pydantic
@@ -6,6 +7,7 @@ import pytest
 import sqlalchemy
 
 import ormar
+from ormar import ModelDefinitionError
 from tests.settings import DATABASE_URL
 
 database = databases.Database(DATABASE_URL, force_rollback=True)
@@ -16,6 +18,11 @@ def time():
     return datetime.datetime.now().time()
 
 
+class MyEnum(Enum):
+    SMALL = 1
+    BIG = 2
+
+
 class Example(ormar.Model):
     class Meta:
         tablename = "example"
@@ -30,6 +37,17 @@ class Example(ormar.Model):
     description: str = ormar.Text(nullable=True)
     value: float = ormar.Float(nullable=True)
     data: pydantic.Json = ormar.JSON(default={})
+    size: MyEnum = ormar.Enum(enum_class=MyEnum, default=MyEnum.SMALL)
+
+
+class EnumExample(ormar.Model):
+    class Meta:
+        tablename = "enum_example"
+        metadata = metadata
+        database = database
+
+    id: int = ormar.Integer(primary_key=True)
+    size: MyEnum = ormar.Enum(enum_class=MyEnum, default=MyEnum.SMALL)
 
 
 @pytest.fixture(autouse=True, scope="module")
@@ -40,6 +58,49 @@ def create_test_database():
     metadata.drop_all(engine)
 
 
+def test_proper_enum_column_type():
+    assert Example.__fields__["size"].type_ == MyEnum
+
+
+def test_accepts_only_proper_enums():
+    class WrongEnum(Enum):
+        A = 1
+        B = 2
+
+    with pytest.raises(pydantic.ValidationError):
+        Example(size=WrongEnum.A)
+
+
+@pytest.mark.asyncio
+async def test_enum_bulk_operations():
+    async with database:
+        examples = [EnumExample(), EnumExample()]
+        await EnumExample.objects.bulk_create(examples)
+
+        check = await EnumExample.objects.all()
+        assert all(x.size == MyEnum.SMALL for x in check)
+
+        for x in check:
+            x.size = MyEnum.BIG
+
+        await EnumExample.objects.bulk_update(check)
+        check2 = await EnumExample.objects.all()
+        assert all(x.size == MyEnum.BIG for x in check2)
+
+
+@pytest.mark.asyncio
+async def test_enum_filter():
+    async with database:
+        examples = [EnumExample(), EnumExample(size=MyEnum.BIG)]
+        await EnumExample.objects.bulk_create(examples)
+
+        check = await EnumExample.objects.all(size=MyEnum.SMALL)
+        assert len(check) == 1
+
+        check = await EnumExample.objects.all(size=MyEnum.BIG)
+        assert len(check) == 1
+
+
 @pytest.mark.asyncio
 async def test_model_crud():
     async with database:
@@ -52,6 +113,13 @@ async def test_model_crud():
         assert example.description is None
         assert example.value is None
         assert example.data == {}
+        assert example.size == MyEnum.SMALL
+
+        await example.update(data={"foo": 123}, value=123.456, size=MyEnum.BIG)
+        await example.load()
+        assert example.value == 123.456
+        assert example.data == {"foo": 123}
+        assert example.size == MyEnum.BIG
 
         await example.update(data={"foo": 123}, value=123.456)
         await example.load()
@@ -59,3 +127,18 @@ async def test_model_crud():
         assert example.data == {"foo": 123}
 
         await example.delete()
+
+
+@pytest.mark.asyncio
+async def test_invalid_enum_field():
+    async with database:
+        with pytest.raises(ModelDefinitionError):
+
+            class Example2(ormar.Model):
+                class Meta:
+                    tablename = "example"
+                    metadata = metadata
+                    database = database
+
+                id: int = ormar.Integer(primary_key=True)
+                size: MyEnum = ormar.Enum(enum_class=[])

tests/test_queries/test_queryset_level_methods.py

@@ -1,4 +1,5 @@
-from typing import List, Optional
+from enum import Enum
+from typing import Optional
 
 import databases
 import pydantic
@@ -19,6 +20,11 @@ database = databases.Database(DATABASE_URL, force_rollback=True)
 metadata = sqlalchemy.MetaData()
 
 
+class MySize(Enum):
+    SMALL = 0
+    BIG = 1
+
+
 class Book(ormar.Model):
     class Meta:
         tablename = "books"
@@ -45,6 +51,7 @@ class ToDo(ormar.Model):
     text: str = ormar.String(max_length=500)
     completed: bool = ormar.Boolean(default=False)
     pairs: pydantic.Json = ormar.JSON(default=[])
+    size = ormar.Enum(enum_class=MySize, default=MySize.SMALL)
 
 
 class Category(ormar.Model):
@@ -77,6 +84,7 @@ class ItemConfig(ormar.Model):
     id: Optional[int] = ormar.Integer(primary_key=True)
     item_id: str = ormar.String(max_length=32, index=True)
     pairs: pydantic.Json = ormar.JSON(default=["2", "3"])
+    size = ormar.Enum(enum_class=MySize, default=MySize.SMALL)
 
 
 class QuerySetCls(QuerySet):
@@ -343,6 +351,7 @@ async def test_bulk_update():
         for todo in todoes:
             todo.text = todo.text + "_1"
             todo.completed = False
+            todo.size = MySize.BIG
 
         await ToDo.objects.bulk_update(todoes)
 
@@ -354,6 +363,7 @@ async def test_bulk_update():
 
         for todo in todoes:
             assert todo.text[-2:] == "_1"
+            assert todo.size == MySize.BIG
 
 
 @pytest.mark.asyncio
@@ -474,3 +484,13 @@ async def test_custom_queryset_cls():
         await Customer(name="test").save()
         c = await Customer.objects.first_or_404(name="test")
         assert c.name == "test"
+
+
+@pytest.mark.asyncio
+async def test_filter_enum():
+    async with database:
+        it = ItemConfig(item_id="test_1")
+        await it.save()
+
+        it = await ItemConfig.objects.filter(size=MySize.SMALL).first()
+        assert it