[BUG] Validation fails in beanie 1.24.0 when dictionary key is an Enum

[juho-workfellow]

Describe the bug

https://github.com/roman-right/beanie/pull/785 introduced a bug here: https://github.com/roman-right/beanie/blob/main/beanie/odm/utils/encoder.py#L134. Now validation fails if dictionary key is an Enum.

To Reproduce

from beanie import Document, init_beanie
from motor.motor_asyncio import AsyncIOMotorClient
from typing import Dict
from enum import Enum
import asyncio


class CategoryEnum(str, Enum):
    CHOCOLATE = "Chocolate"
    CANDY = "Candy"

class Product(Document):
    name: str
    price: float
    category: Dict[CategoryEnum, str]

async def example():
    client = AsyncIOMotorClient(<mongodb_connection>)
    await init_beanie(database=client["test_db"], document_models=[Product])

    chocolate = {CategoryEnum.CHOCOLATE: "description of chocolate"}
    tonybar = Product(name="Tony's", price=1.1, category=chocolate)
    await tonybar.insert()

    product = await Product.find_one()
    print(product)


if __name__ == "__main__":
    asyncio.run(example())

Expected behavior With beanie 1.23.6 the script prints this on the command line:

id=ObjectId('659fc96a59f26f162d6cf1dc') revision_id=None name="Tony's" price=1.1 category={<CategoryEnum.CHOCOLATE: 'Chocolate'>: 'description of chocolate'}

With beanie 1.24.0:

Traceback (most recent call last):
  File "C:some_path\beanie_bug.py", line 36, in <module>
    asyncio.run(example())
  File "C:\Users\JuhoSalminen\AppData\Local\Programs\Python\Python310\lib\asyncio\runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "C:\Users\JuhoSalminen\AppData\Local\Programs\Python\Python310\lib\asyncio\base_events.py", line 649, in run_until_complete
    return future.result()
  File "C:some_path\beanie_bug.py", line 31, in example
    product = await Product.find_one()
  File "C:some_path\.venv\lib\site-packages\beanie\odm\queries\find.py", line 1023, in __await__
    FindQueryResultType, parse_obj(self.projection_model, document)
  File "C:some_path\.venv\lib\site-packages\beanie\odm\utils\parsing.py", line 104, in parse_obj
    result = parse_model(model, data)
  File "C:some_path\.venv\lib\site-packages\beanie\odm\utils\pydantic.py", line 39, in parse_model
    return model_type.parse_obj(data)
  File "pydantic\main.py", line 526, in pydantic.main.BaseModel.parse_obj
  File "C:some_path\.venv\lib\site-packages\beanie\odm\documents.py", line 182, in __init__
    super(Document, self).__init__(*args, **kwargs)
  File "pydantic\main.py", line 341, in pydantic.main.BaseModel.__init__
pydantic.error_wrappers.ValidationError: 1 validation error for Product
category -> __key__
  value is not a valid enumeration member; permitted: 'Chocolate', 'Candy' (type=type_error.enum; enum_values=[<CategoryEnum.CHOCOLATE: 'Chocolate'>, <CategoryEnum.CANDY: 'Candy'>])

Replacing line 134 in beanie/odm/utils/encoder.py with return {key: self.encode(value) for key, value in obj.items()} fixes the issue.

[kheyer]

Just ran into this. Resolved by downgrading from 1.24.0 to 1.23.6.

[roman-right]

Hi! Thank you for the catch! I'll pick it up on the next bug fixing session