Skip to content

Instantly share code, notes, and snippets.

@hjwp
Created October 27, 2020 14:03
Show Gist options
  • Save hjwp/405f04802ea558f042728ec5edbb4e62 to your computer and use it in GitHub Desktop.
Save hjwp/405f04802ea558f042728ec5edbb4e62 to your computer and use it in GitHub Desktop.
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]
@MartinCura
Copy link

StrEnum didn't exist back when you wrote this (new in 3.11), did they base it on your (and mine) frustration? https://docs.python.org/3/library/enum.html#enum.StrEnum

@hjwp
Copy link
Author

hjwp commented May 9, 2023

ha! must add a little addeundum to the blog post :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment