Search code examples
python-3.xenumspython-dataclasses

Utilizing a dataclass and enum together


I'm trying to write a class that contains both behavior and static instances of the objects it defines, in doing this I'm attempting to use dataclass(frozen=True) and enum.

@dataclass(frozen=True)
class Foo(Enum):
    a: int
    b: float

    FOO1 = Foo(1, 1.0)
    FOO2 = Foo(2, 0.5)

An obvious complication of this approach is that you cannot define a class level variable in terms of the class itself, so I'm forced to utilize something like this:


@dataclass(frozen=True)
class _Foo:
    a: int
    b: float


class Foo(Enum):
    FOO1 = _Foo(1, 1.0)
    FOO2 = _Foo(2, 0.5)

Doing this however causes issues with type hinting as the code thinks that FOO1 and FOO2 are instances of Foos rather than Foo. So okay, that could be a simple fix, so then I go with:


@dataclass(frozen=True)
class _Foo:
    a: int
    b: float


@dataclass(frozen=True)
class Foo(_Foo, Enum):
    FOO1 = _Foo(1, 1.0)
    FOO2 = _Foo(2, 0.5)

This now causes the program to error with TypeError: _value_ not set in __new__, unable to create it, and I am now at a loss. Is there a way to blend a dataclass with an enum effectively, or is there a better approach to this?


Solution

  • I don't know if this works well with typing (please post a comment letting me know), but, as of Python3.11, Enum is optimized to work with dataclass using the following pattern:

    @dataclass(frozen=True)
    class Foo:
        a: int
        b: float
    
    class FooEnum(Foo, Enum):
        FOO1 = 1, 1.0
        FOO2 = 2, 0.5
    

    and in use:

    >>> FooEnum.FOO1
    <FooEnum.FOO1: a=1, b=1.0>
    
    >>> FooEnum.FOO1.value
    Foo(a=1, b=1.0)