Python enum with custom class with mandatory keyword arguments

Question

I am trying to create an enum with custom class which has mandatory positional arguments.

class Coord:
  def __init__(self, *, x, y):
    self._x = x
    self._y = y

I would like to make several enums with this class. Somthing like:

class IntCoord(Coord, Enum):
  ONE_ONE = Coord(x=1, y=1)

class FloatCoord(Coord, Enum):
  ONE_ONE = Coord(x=0.1, y=0.1)

So that I can use them in other class like:

class System(Generic[T]):
  @classmethod
  def get_one_one(cls) -> T:
    return NotImplemented

class IntSystem(System[IntCoord]):
  @classmethod
  def get_one_one(cls) -> IntCoord:
    return IntCoord.ONE_ONE

class FloatSystem(System[FloatCoord]):
  @classmethod
  def get_one_one(cls) -> FloatCoord:
    return FloatCoord.ONE_ONE

After that, I want to compare the enums:

IntCoord(x=1, y=1) is IntSystem.get_one_one()

I am having a hard time defining the Coord class. It seems it is not possible to instantiate Coord enum class by Coord(x=1, y=1)

Is there a good way to achieve this with enum? Or better reasonable implementation?

Full script is as follows:

from enum import Enum
from typing import Generic, TypeVar


class Coord:
    def __init__(self, *, x, y):
        self._x = x
        self._y = y


class IntCoord(Coord, Enum):
    ONE_ONE = Coord(x=1, y=1)


class FloatCoord(Coord, Enum):
    ONE_ONE = Coord(x=0.1, y=0.1)


T = TypeVar('T', bound=Coord)


class System(Generic[T]):
    @classmethod
    def get_one_one(cls) -> T:
        return NotImplemented


class IntSystem(System[IntCoord]):
    @classmethod
    def get_one_one(cls) -> IntCoord:
        return IntCoord.ONE_ONE


class FloatSystem(System[FloatCoord]):
    @classmethod
    def get_one_one(cls) -> FloatCoord:
        return FloatCoord.ONE_ONE


if __name__ == '__main__':
    int_one_one = IntSystem.get_one_one()
    float_one_one = FloatSystem.get_one_one()

    print(IntCoord(x=1, y=1) is int_one_one)
    print(FloatCoord(x=0.1, y=0.1) is float_one_one)

Errors are:

Traceback (most recent call last):
  File "test.py", line 11, in 
    class IntCoord(Coord, Enum):
  File ".pyenv/versions/3.9.0/lib/python3.9/enum.py", line 228, in __new__
    enum_member.__init__(*args)
TypeError: __init__() takes 1 positional argument but 2 were given

user_208880 · Answer

The problem you are having with Coord and Enum is that you have said your enum is also a Coord, and a Coord takes two arguments -- but you are only giving your enum member ONE_ONE a single argument... a Coord.

There are two ways around this problem:

don't use a Coord in your enum member (and lose your keyword (not positional) arguments):

    class Coord:
        def __init__(self, x, y):
            self._x = x
            self._y = y

    class IntCoord(Coord, Enum):

        ONE_ONE = 1, 1

don't use Coord in your enum header (and keep your keyword arguments, but add propertys so that blah._x and blah._y still work]:

    class Coord:
        def __init__(self, *, x, y):
            self._x = x
            self._y = y

    class IntCoord(Enum):
        #
        ONE_ONE = Coord(x=1, y=1)
        #
        @property
        def _x(self):
            return self.value._x
        #
        @property
        def _y(self):
            return self.value._y

It does get easier if you use the aenum¹ library instead:

from aenum import Enum, enum

class Coord:
    def __init__(self, *, x, y):
        self._x = x
        self._y = y

class IntCoord(Coord, Enum):
    #
    ONE_ONE = enum(x=1, y=1)

and in use:

>>> IntCoord.ONE_ONE


>>> IntCoord.ONE_ONE._y
1

Note that in all the above cases, retrieving IntCoord.ONE_ONE by value is a little trickier than usual:

IntCoord((1, 1))
IntCoord(Coord(x=1, y=1))
IntCoord(enum(x=1, y=1))

--

¹ Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

Python enum with custom class with mandatory keyword arguments

1 Answers