Using different Enums (or Literal) on same data model, for different fields?

[impaktor]

I have a solution (see code below) to a problem, but I'm thinking there is perhaps a better way? I've half-contemplated using dict, or Literal, but solution below uses Enum.

The problem: I have a large amount of fields that are all of key-value pairs (must be named @V for value, and @N for name), that I define in an enum MyPairs in the code below, but which key/values depends on the field. E.g. the field sex takes key-values: {unknown: 0, man: 1, woman: 2}, while compression field takes key-values {gz: "gzip", df: "defalate"}, etc. I have close to 100 different enums like that (spanning total +1000 lines) for these fields, that define which key-value pairs are valid for a specific field of various models.

For this small example, the goal is to get it to dictionary of type:

{'compression': {'@N': 'Gzip', '@V': 'GZ'},
 'sex':         {'@N': 'Woman', '@V': '2'}}

(If I have the dict above, it's trivial to then use xmltodict.unparse() to convert the @ fields as XML attributes, below

<sex V="1" N="Man"></sex>
<compression V="DF" N="Deflate"></compression>

)

Code

The code below gives the desired output, but to get there, I have to define custom @classmethod for the factory of every field (of which I have many). So wondering if there's a way to achieve same output, but with less code?

from enum import Enum

import xmltodict
from polyfactory.factories.pydantic_factory import ModelFactory
from pydantic import BaseModel, Field
from pprint import pprint


class MyPairs(BaseModel):
    """For holding key-value pairs: V (value) + N (name)"""

    V: str = Field(alias="@V", description="Value")
    N: str = Field(alias="@N", description="Name")

    class Config:
        """Assign values using field name, rather than alias names"""
        populate_by_name = True


class Sex(Enum):
    Unknown = 0
    Man = 1
    Woman = 2
    Not_applicable = 9


class Compression(Enum):
    Deflate = "DF"
    Gzip = "GZ"
    Compress = "Z"


class Person(BaseModel):
    name: str
    sex: MyPairs                  # ideally specify it takes values defined in "Sex" Enum
    compression: MyPairs          # ideally specify it takes values defined in "Compression" Enum


class PersonFactory(ModelFactory[Person]):

    @classmethod
    def sex(cls) -> MyPairs:
        # Need to unpack the enum, to populate the MyPairs object
        e = cls.__random__.choice(list(Sex))
        return MyPairs(V=str(e.value), N=e.name)

    @classmethod
    def compression(cls) -> MyPairs:
        e = cls.__random__.choice(list(Compression))
        return MyPairs(V=e.value, N=e.name)


# The "Person's" Pet will have name that's random string, rather than generated by "Faker" module.
fake_person = PersonFactory.build()
print("\nPERSON:")
print(fake_person)
print("\n")

d_fake = {"root": fake_person.dict(by_alias=True)}
pprint(d_fake)
xml_fake = xmltodict.unparse(d_fake, pretty=True).replace("\t", "  ")
print(xml_fake)

Alternative

Why not use the Enum directly, instead of unpacking it in a @classmethod, might be the obvious question?
like below:

class Person(BaseModel):
    name: str
    sex: MyPairs
    compression: Compression

Well, then output is:

{'compression': <Compression.Compress: 'Z'>}

Rather then the desired

{'compression': {'@N': 'Compress', '@V': 'Z'},

I've experimented with overloading various dunder methods of the Enum class, such that it wraps the key values in a MyPairs data class (semantically similar to what's done in @classmethod), but haven't converged on anything that works, though I thought this seemed promising.

Possibly I should overload the function that converts the pydantic data model to dict?

[adhtruong]

I think this is expected behaviour from polyfactory. For first idea, without type information on specific type so need to specify this explicitly.

For alternative, assuming pydantic 2 something like this might work. Note this is incomplete if you want to load data

from enum import Enum
from typing import TypeVar

from typing_extensions import Annotated

from pydantic import BaseModel, PlainSerializer

T = TypeVar("T", bound=Enum)
EnumSerialiser = Annotated[T, PlainSerializer(lambda x: {"@N": x.name, "@V": str(x.value)})]


class MyEnum(Enum):
    A = 1
    B = 2


class MyModel(BaseModel):
    field: EnumSerialiser[MyEnum]


print(MyModel(field=MyEnum.A).model_dump())

[impaktor]

That's very interesting! Much better solution than I had in mind, thanks!