CODEWITHSUNDEEP
pythonpython-2.7python-3.xcastingenums
AtPython
AtPython

Reputation: 193

Cast int to enum in Python for my program?

How can a int be cast to an enum in python?

Upvotes: 6

Views: 8171

Answers (3)

crifan
crifan

Reputation: 14338

For here latest Python 3.7 and some Enum:

from enum import Enum

class BatteryState(Enum):
    Unknown = 0
    Unplugged = 1
    Charging = 2
    Full = 3

Convert Int to Enum

just pass int to Enum

batteryStateInt = 2
curBattryState = BatteryState(batteryStateInt)
print("curBattryState=%s" % curBattryState)

can got what you expected Enum:

curBattryState=BatteryState.Charging

Extra

get Enum string expression can via name

curBattryStateName = curBattryState.name
print("curBattryStateName=%s" % curBattryStateName)

can got expected: Charging

Upvotes: 0

Ethan Furman
Ethan Furman

Reputation: 69288

You can use the built-in Enum (Python 3.4+), the enum34 backport, or, for more advanced needs (which this is not), the new aenum library.

If you use IntEnum:

class RGB(IntEnum):
    red = 1
    green = 2
    blue = 3

If you have an int and want the matching Enum member:

>>> c = 2
>>> c = RGB(c)
>>> c
<RGB.green: 2>

Once you have an IntEnum member, it is already an int:

>>> type(c)
<enum 'RGB'>
>>> isinstance(c, int)
True

The downside to IntEnum is that every IntEnum member will compare equal to every other IntEnum member that has the same value:

class Fruit(IntEnum):
    banana = 1

>>> Fruit.banana == Color.red
True

If you want to keep your Enums separate, but don't want to lose the intability of an IntEnum you could do:

class RGB(Enum):
    red = 1
    green = 2
    blue = 3
    def __int__(self):
        return self.value

Lookup works the same:

>>> c = 2
>>> c = RGB(c)
>>> c
<RGB.green: 2>

But members are no longer ints:

>>> type(c)
<enum 'RGB'>
>>> isinstance(c, int)
False

So you do have to cast it (or access the value attribute):

>>> int(c)
2
>>> c.value
2

And we no longer have the problem if Fruit.banana being equal to RGB.red:

>>> Fruit(1) == RGB(1)
False

If you are using Python 3+ there are some cool things you can do with aenum, such as:

import aenum

class RGB(aenum.Enum, start=1):
    red
    green
    blue
    def __int__(self):
        return self.value

which results in the same class as the last Enum example.

Upvotes: 2

lemonhead
lemonhead

Reputation: 5518

If you want the flexibility to convert between int and an enum, you can use enum.IntEnum

import enum

class Color(enum.IntEnum):
    green = 1
    blue = 2
    red = 3
    yellow = 4

color_code = 4
# cast to enum
color = Color(color_code)

# cast back to int
color_code = int(color)

Note: If you are using python<3.4, enum has been backported, but you will need to install it, e.g. via pip install enum

More on enums in python - https://docs.python.org/3/library/enum.html

Upvotes: 11

Related Questions