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