Make a (typed) python enum with auto value based on name and also options

user_8031815 0
Viewed 342

In python you can define a typed enum with auto values:

import enum
from enum import auto

class Ordinals(enum.IntEnum):
  FIRST = auto()
  SECOND = auto()
  THIRD = auto()

Ordinals.FIRST == 1 #=> True

And you can also define a typed enum with arguments: (Example from docs):

class Coordinate(bytes, Enum):
     """
     Coordinate with binary codes that can be indexed by the int code.
     """
     def __new__(cls, value, label, unit):
         obj = bytes.__new__(cls, [value])
         obj._value_ = value
         obj.label = label
         obj.unit = unit
         return obj
     PX = (0, 'P.X', 'km')
     PY = (1, 'P.Y', 'km')
     VX = (2, 'V.X', 'km/s')
     VY = (3, 'V.Y', 'km/s')

Coordinate.PX._name_ #=> 'PX'
Coordinate.PX == bytes([0]) #=> True
Coordinate.PX.label #=> 'P.X'

Furthermore, you can define an enum that looks up a value by the name:

class LetterOrdinals(enum.IntEnum):
  def _generate_next_value_(name, *_):
    return ord(name)
  A = auto()
  B = auto()

LetterOrdinals.A == 65 #=> True

Now say I want to combine all three and have a typed enum that generates a value from the name and allows for arguments to __new__/__init__. Is this possible? The first problem I run into is that auto no longer works on the value. it seems like if auto isn't the entire value, then it skips generating the value. Conversely, if it doesn't generate the value, the name is not available to __new__ or __init__. Here's an example that doesn't work (and I've tried variants on this):

class Coordinate(bytes, Enum):
     def _generate_next_value_(name, *_):
         return [ord(letter) for letter in value]
     def __new__(cls, value, label, unit):
         obj = bytes.__new__(cls, [*value])
         obj._value_ = value
         obj.label = label
         obj.unit = unit
         return obj
     PX = (auto(), 'P.X', 'km')
     PY = (auto(), 'P.Y', 'km')
     VX = (auto(), 'V.X', 'km/s')
     VY = ([2,2], 'V.Y', 'km/s')
# never generates the auto value, causes errors

So is there a good way around this?

enums python
1 Answers

You can write your own item definition method in the class and have it return the object created by auto(). The Enum class does some special processing with the enum.auto objects so you need to temporarily store the other attributes somewhere until the __init__() is called and retrieve them for actual instance initialization:

from enum import Enum,auto
class Coordinate: args = [] # placeholder for auto/init arguments
class Coordinate(Enum):

     def _generate_next_value_(name, *_):
        # Here, you can use label,unit = Coordinate.args[-1] 
        # if your value is based on other arguments
        return ord(name[0])*100+ord(name[1])

     def __init__(obj,value):
         obj.label,obj.unit = Coordinate.args.pop(0)

     def item(*args):
         Coordinate.args.append(args)
         return auto()
        
     PX = item('P.X', 'km')
     PY = item('P.Y', 'km')
     VX = item('V.X', 'km/s')
     VY = item('V.Y', 'km/s')

Output:

print(Coordinate['VY'])        # Coordinate.PY
print(Coordinate['VY'].value)  # 8689
print(Coordinate['VY'].label)  # P.Y
print(Coordinate['VY'].unit)   # km
print(Coordinate(8689))        # Coordinate.PY
print(Coordinate(8689).name)   # PY
user_52375600
Related