Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

Tags:

python

enums

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?

like image 660
Garrett Motzner Avatar asked Jun 30 '26 21:06

Garrett Motzner


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
like image 177
Alain T. Avatar answered Jul 03 '26 10:07

Alain T.



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!