Re: [Python-ideas] constant/enum type in stdlib

Feb. 1, 2013

      Wow - this ended up being more difficult that I'd anticipated. Ensuring
that decorators work in Michael Foord-inspired enums, etc mandated a fairly
big redesign. As it was you couldn't do something like:

class A():
    a = 1

class B(Enum):
    B = A.a

because when the name 'A' was resolved in the definition of class B, it was
returning an int (the next enum value).

The end result may well still have some holes, but it's looking pretty good
to me.

Enums are constructed like:

class Color(Enum):
    RED, GREEN, BLUE
    CYAN = 10
    MAGENTA
    YELLOW
    BLACK

where the assigned enum starts a new count i.e. the above is 0, 1, 2, 10,
11, 12, 13. Arbitrary attributes may be assigned and not contribute to the
enumeration so long as the object being assigned does not implement
__index__ (if it does, it creates a discontiguous enumeration).

#!/usr/bin/env python3

import builtins
import collections
import operator

class EnumValue(int):
    def __new__(cls, key, value):
        e = super().__new__(cls, value)
        super().__setattr__(e, 'key', key)
        super().__setattr__(e, 'owner', None)
        return e

    def __setattr__(self, key, value):
        raise TypeError("Cannot set attribute of type %r" % (type(self),))

    def __str__(self):
        return "%s.%s" % (self.owner.__name__, self.key)

    def __repr__(self):
        if self.owner is not None:
            return "<%s '%s.%s': %d>" % (self.__qualname__,
self.owner.__qualname__, self.key, int(self))

        return "<%s '%s': %d>" % (self.__qualname__, self.key, int(self))

class EnumProxy(object):
    def __init__(self, key, value=None):
        self.key = key
        self.values = [value]
        self.used = False

    def __repr__(self):
        return "<%s '%s': %s>" % (self.__qualname__, self.key, self.values)

    def _get(self, used=None):
        if used:
            self.used = True

        try:
            return locals()[self.key]
        except KeyError:
            try:
                return globals()[self.key]
            except KeyError:
                try:
                    return getattr(builtins, self.key)
                except KeyError:
                    raise NameError(self.key, self.values)

    def __call__(self, *p, **kw):
        return self._get(True)(*p, **kw)

    def __getattr__(self, name):
        return getattr(self._get(True), name)

class EnumValues(collections.OrderedDict):
    def __init__(self):
        super().__init__()
        self.sealed = False

    def __getitem__(self, key):
        try:
            obj = super().__getitem__(key)

            if not self.sealed and isinstance(obj, EnumProxy):
                obj.values.append(None)

            return obj

        except KeyError:
            # Don't do anything with __dunder__ attributes
            if key[:2] == '__' and key[-2:] == '__':
                raise

            proxy = EnumProxy(key, None)
            super().__setitem__(key, proxy)
            return proxy

    def __setitem__(self, key, value):
        if key[:2] == '__' and key[-2:] == '__':
            return super().__setitem__(key, value)

        try:
            if isinstance(value, EnumProxy):
                value = value._get(True)

            elif not isinstance(value, EnumValue):
                value = operator.index(value)

        except TypeError:
            return super().__setitem__(key, value)

        try:
            o = super().__getitem__(key)

            if isinstance(o, EnumProxy):
                o.values.append(value)

        except KeyError:
            if isinstance(value, EnumProxy):
                int.__setattr__(value, 'key', key)
            else:
                value = EnumProxy(key, value)

            super().__setitem__(value.key, value)

class EnumMeta(type):

    @classmethod
    def __prepare__(metacls, name, bases):
        return EnumValues()

    def __new__(cls, name, bases, classdict):
        classdict.sealed = True

        del_list = []

        for v in classdict.values():
            if isinstance(v, EnumProxy) and v.used:
                del_list.append(v)

        for v in del_list:
            del classdict[v.key]

        result = type.__new__(cls, name, bases, dict(classdict))
        value = 0
        keys = {}
        values = {}

        for v in classdict.values():
            if isinstance(v, EnumProxy) and not v.used:
                if len(v.values) > 1:
                    raise AttributeError("Duplicate enum key '%s.%s'" %
(result.__qualname__, v.key,))

                if v.values[0] is not None:
                    value = v.values[0]

                if isinstance(value, EnumValue):

                    if (value.key is not None) and (value.key != v.key):
                        raise AttributeError("Assigned enum value to
non-matching key '%s': %r" % (v.key, value))

                    if value.owner is not None:
                        raise AttributeError("Assigned owned enum value to
key '%s': %r" % (v.key, value))

                    int.__setattr__(value, 'key', v.key)
                    v = value

                else:
                    v = EnumValue(v.key, value)

                setattr(result, v.key, v)
                value += 1

            if isinstance(v, EnumValue):
                int.__setattr__(v, 'owner', result)

                if v in values:
                    raise AttributeError("Duplicate enum value %d for keys:
'%s' and '%s'" % (int(v), values[v].key, v.key))

                keys[v.key] = v
                values[v] = v

        enum = sorted(values)

        result._key_to_enum = collections.OrderedDict()
        result._value_to_enum = values

        for e in enum:
            result._key_to_enum[e.key] = e

        return result

    def __getitem__(self, key):
        try:
            key = operator.index(key)
        except TypeError:
            return self._key_to_enum[key]
        else:
            return self._value_to_enum[key]

    def _items(self):
        return self._key_to_enum.items()

    def _keys(self):
        return self._key_to_enum.keys()

    def _values(self):
        return self._key_to_enum.values()

    def items(self):
        return self._items()

    def keys(self):
        return self._keys()

    def values(self):
        return self._values()

    def __iter__(self):
        return iter(self.values())

    def __repr__(self):
        r = super().__repr__()
        r = ['<enum', r[6:-1], ' ']
        r.append(str(self))
        r.append('>')
        return ''.join(r)

    def __str__(self):
        s = ['{']

        for k, v in self.items():
            if s[-1][-1:] != '{':
                s.append(', ')

            s.extend([k, ':', str(int(v))])

        s.append('}')
        return ''.join(s)

class Enum(metaclass=EnumMeta):
    def __getitem__(self, key):
        cls = type(self)
        return type(cls).__getitem__(cls, key)

    @classmethod
    def items(cls):
        return cls._items()

    @classmethod
    def keys(cls):
        return cls._keys()

    @classmethod
    def values(cls):
        return cls._values()

    def __iter__(self):
        return iter(self.values())

    def __repr__(self):
        r = super().__repr__()
        r = r.replace('object at 0x', 'enum at 0x')
        r = [r[:-1], ' ']
        r.append(str(self))
        r.append('>')
        return ''.join(r)

    def __str__(self):
        return str(type(self))

if __name__ == '__main__':

    class Color(Enum):
        RED, GREEN, BLUE
        ORANGE = "orange"
        CYAN = 10
        MAGENTA
        YELLOW
        BLACK

    import unittest

    class TestEnum(unittest.TestCase):

        EXPECTED_KEY_ORDER = ('RED', 'GREEN', 'BLUE', 'CYAN', 'MAGENTA',
'YELLOW', 'BLACK')
        EXPECTED_INT_VALUE_ORDER = (0, 1, 2, 10, 11, 12, 13)
        EXPECTED_ENUM_VALUE_ORDER = (Color.RED, Color.GREEN, Color.BLUE,
Color.CYAN, Color.MAGENTA, Color.YELLOW, Color.BLACK)
        EXPECTED_ITEMS_ORDER = tuple(zip(EXPECTED_KEY_ORDER,
EXPECTED_ENUM_VALUE_ORDER))

        def test_type(self):
            self.assertIsInstance(Color.RED, int)

        def test_class_enum_values(self):
            self.assertEqual(0, Color.RED)
            self.assertEqual(1, Color.GREEN)
            self.assertEqual(2, Color.BLUE)
            self.assertEqual(10, Color.CYAN)
            self.assertEqual(11, Color.MAGENTA)
            self.assertEqual(12, Color.YELLOW)
            self.assertEqual(13, Color.BLACK)
            self.assertEqual("orange", Color.ORANGE)

        def test_instance_enum_values(self):
            e = Color()
            self.assertIs(Color.RED, e.RED)
            self.assertIs(Color.GREEN, e.GREEN)
            self.assertIs(Color.BLUE, e.BLUE)
            self.assertIs(Color.CYAN, e.CYAN)
            self.assertIs(Color.MAGENTA, e.MAGENTA)
            self.assertIs(Color.YELLOW, e.YELLOW)
            self.assertIs(Color.BLACK, e.BLACK)
            self.assertIs(Color.ORANGE, e.ORANGE)

        def test_class_indexing(self):
            self.assertIs(Color.CYAN, Color['CYAN'])
            self.assertIs(Color.CYAN, Color[10])

        def test_instance_indexing(self):
            e = Color()
            self.assertIs(Color.CYAN, e['CYAN'])
            self.assertIs(Color.CYAN, e[10])

        def test_class_keys(self):
            self.assertEqual(self.EXPECTED_KEY_ORDER, tuple(Color.keys()))

        def test_instance_keys(self):
            self.assertEqual(tuple(Color.keys()), tuple(Color().keys()))

        def test_class_values(self):
            self.assertEqual(self.EXPECTED_INT_VALUE_ORDER,
tuple(Color.values()))
            self.assertEqual(self.EXPECTED_ENUM_VALUE_ORDER,
tuple(Color.values()))

        def test_instance_values(self):
            self.assertEqual(tuple(Color.values()), tuple(Color().values()))

        def test_class_items(self):
            self.assertEqual(self.EXPECTED_ITEMS_ORDER,
tuple(Color.items()))

        def test_instance_items(self):
            self.assertEqual(tuple(Color.items()), tuple(Color().items()))

        def test_owner(self):
            for e in Color:
                self.assertIs(e.owner, Color)

        def test_class_str(self):
            s = str(Color)

            for e in Color:
                self.assertIn('%s:%d' % (e.key, int(e)), s)

        def test_instance_str(self):
            self.assertEqual(str(Color), str(Color()))

        def test_class_repr(self):
            r = repr(Color)
            self.assertIn(Color.__qualname__, r)
            self.assertIn(str(Color), r)

        def test_instance_repr(self):
            e = Color()
            r = repr(e)
            self.assertIn(Color.__qualname__, r)
            self.assertIn('at 0x', r)
            self.assertIn(str(Color()), r)

        def _create_duplicate_key(self):
            class DuplicateKey(Enum):
                KEY,
                KEY

        def test_duplicate_key(self):
            self.assertRaises(AttributeError, self._create_duplicate_key)

        def _create_duplicate_value(self):
            class DuplicateValue(Enum):
                KEY1,
                KEY2 = 0

        def test_duplicate_value(self):
            self.assertRaises(AttributeError, self._create_duplicate_value)

        def _assign_wrong_key(self):
            class WrongKey(Enum):
                KEY1 = EnumValue('KEY2', 0)

        def test_wrong_key(self):
            self.assertRaises(AttributeError, self._assign_wrong_key)

        def test_unnamed_key1(self):
            class UnnamedKey(Enum):
                KEY1 = EnumValue(None, 5)

            self.assertEqual(UnnamedKey.KEY1, 5)
            self.assertIs(UnnamedKey, UnnamedKey.KEY1.owner)

        def test_unnamed_key1(self):
            unnamed = EnumValue(None, 5)

            class UnnamedKey(Enum):
                KEY1 = unnamed

            self.assertEqual(UnnamedKey.KEY1, 5)
            self.assertIs(UnnamedKey.KEY1, unnamed)
            self.assertIs(UnnamedKey, UnnamedKey.KEY1.owner)

        def _assign_wrong_owner(self):
            class WrongOwner(Enum):
                KEY1 = Color.RED

        def test_wrong_owner(self):
            self.assertRaises(AttributeError, self._assign_wrong_owner)

    unittest.main()

Tim Delaney