Better string enums

test_enums.py

import random
from enum import Enum, EnumMeta, IntEnum


class GoatEnumMeta(EnumMeta):
    def __getitem__(cls, item):
        if isinstance(item, int):
            item = cls._member_names_[item]
        return super().__getitem__(item)


class BRAIN(str, Enum, metaclass=GoatEnumMeta):
    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, metaclass=GoatEnumMeta):
    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]