pythonenumspydantic

Is an instance of my Pydantic model also in my Enum?

In my situation I get a response from an API which follows the structure of MyGenericModel. After I parse the response using MyGenericModel, I want to check if it belongs to my set defined in MyEnum. If it does I want to continue with it, if not.. then we proceed further.

here is a small simplified example of my code:

from pydantic import BaseModel
from enum import Enum


class MyGenericModel(BaseModel):
    value_one: str
    value_two: str
    
    def __eq__(self, other):
        return self.model_dump() == other.model_dump()


class MySpecificModel(MyGenericModel):
    value_one: str = "a"
    value_two: str = "b"


class MyEnum(Enum):
    my_model = MySpecificModel()


# This is true
MyGenericModel(value_one="a", value_two="b")  == MySpecificModel()

# This fails - and I want to be true
MyGenericModel(value_one="a", value_two="b") in MyEnum

How to solve this issue?

Solution

  • The problem here was that the 'in' operator was not properly defined in this situation.

    The 'in' operator calls the '_contains_' magic method, but this should be defined in the class that creates the Enum class!
    In order to do that, we need to use the EnumMeta class, inherit from that, and define a new class. This new class will be used as metaclass in MyEnum:

    from enum import EnumMeta

class MetaEnum(EnumMeta):
    def __contains__(cls, item):
        return any([item == x.value for x in cls.__members__.values()])


class MyEnum(Enum, metaclass=MetaEnum):
    my_model = MySpecificModel()

    Using this to redefine MyEnum allows us to check if a generic class is in the Enum, when that is initialized with attributes that make it essentially identical to the specific class:

    # This is true
MyGenericModel(value_one="a", value_two="b")  == MySpecificModel()
# This is now also True
MyGenericModel(value_one="a", value_two="b") in MyEnum

    I have learned a lot from this post about metaclasses:
    What are metaclasses in Python?