Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Better string enums
import random
from enum import Enum, IntEnum
class BRAIN(str, Enum):
SMALL = "small"
MEDIUM = "medium"
GALAXY = "galaxy"
def __str__(self) -> str:
return str.__str__(self)
foo = BRAIN.SMALL # type: str
bar = BRAIN.SMALL # type: BRAIN
def test1():
assert BRAIN.SMALL == "small"
def test2():
assert str(BRAIN.SMALL) == "small"
def test3():
assert BRAIN.SMALL.value == "small"
def test4():
assert [thing for thing in BRAIN] == ["small", "medium", "galaxy"]
def test5():
assert list(BRAIN) == ["small", "medium", "galaxy"]
def test6():
assert [thing.value for thing in BRAIN] == ["small", "medium", "galaxy"]
def test7():
assert random.choice(BRAIN) in ["small", "medium", "galaxy"]
def test8():
assert random.choice(list(BRAIN)) in ["small", "medium", "galaxy"]
class IBRAIN(IntEnum):
SMALL = 1
MEDIUM = 2
GALAXY = 3
def test_ints_1():
assert IBRAIN.SMALL == 1
def test_ints_2():
assert int(IBRAIN.SMALL) == 1
def test_ints_3():
assert IBRAIN.SMALL.value == 1
def test_ints_4():
assert [thing for thing in IBRAIN] == [1, 2, 3]
def test_ints_5():
assert list(IBRAIN) == [1, 2, 3]
def test_ints_6():
assert [thing.value for thing in IBRAIN] == [1, 2, 3]
def test_ints_7():
assert random.choice(IBRAIN) in [1, 2, 3]
def test_ints_8():
assert random.choice(list(IBRAIN)) in [1, 2, 3]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.