Python Enum prevent invalid attribute assignment

Question

When I use the Functional API to create an Enum, I get back an Enum object that allows arbitrary assignment (i.e. it has a __dict__):

e = enum.Enum('Things',[('foo',1),('bar',2)])
e.baz = 3

The item does not appear in the list:

list(e)
[<foo.foo: 1>, <foo.bar: 2>]

But it can still be referenced:

if thing == e.baz: ...

Now while it seems unlikely to ever occur, one of the reasons I want to use an Enum is to prevent spelling mistakes and string literals, and for those things to be caught when a module is imported or as early as possible.

Is there a way to dynamically build an Enum that behaves more like a __slots__ object that does not allow arbitrary attributes to be assigned?

Answer 1

To make an enum class fully 'read-only', all that is required is a meta class that uses the __setattr__ hook that prevents all attribute assignments. Because the metaclass is attached to the class after it is created, there is no issue with assigning the proper enumerated values.

Like Ethan's answer, I'm using the EnumMeta class as a base for the custom metaclass:

from enum import EnumMeta, Enum

class FrozenEnumMeta(EnumMeta):
    "Enum metaclass that freezes an enum entirely"
    def __new__(mcls, name, bases, classdict):
        classdict['__frozenenummeta_creating_class__'] = True
        enum = super().__new__(mcls, name, bases, classdict)
        del enum.__frozenenummeta_creating_class__
        return enum

    def __call__(cls, value, names=None, *, module=None, **kwargs):
        if names is None:  # simple value lookup
            return cls.__new__(cls, value)
        enum = Enum._create_(value, names, module=module, **kwargs)
        enum.__class__ = type(cls)
        return enum

    def __setattr__(cls, name, value):
        members = cls.__dict__.get('_member_map_', {})
        if hasattr(cls, '__frozenenummeta_creating_class__') or name in members:
            return super().__setattr__(name, value)
        if hasattr(cls, name):
            msg = "{!r} object attribute {!r} is read-only"
        else:
            msg = "{!r} object has no attribute {!r}"
        raise AttributeError(msg.format(cls.__name__, name))

    def __delattr__(cls, name):
        members = cls.__dict__.get('_member_map_', {})
        if hasattr(cls, '__frozenenummeta_creating_class__') or name in members:
            return super().__delattr__(name)
        if hasattr(cls, name):
            msg = "{!r} object attribute {!r} is read-only"
        else:
            msg = "{!r} object has no attribute {!r}"
        raise AttributeError(msg.format(cls.__name__, name))

class FrozenEnum(Enum, metaclass=FrozenEnumMeta):
    pass

The above distinguishes between attributes that are already available and new attributes, for ease of diagnosing. It also blocks attribute deletion, which is probably just as important!

It also provides both the metaclass and a FrozenEnum base class for enumerations; use this instead of Enum.

To freeze a sample Color enumeration:

>>> class Color(FrozenEnum):
...     red = 1
...     green = 2
...     blue = 3
...
>>> list(Color)
[<Color.red: 1>, <Color.green: 2>, <Color.blue: 3>]
>>> Color.foo = 'bar'
Traceback (most recent call last):
    # ...
AttributeError: 'Color' object has no attribute 'foo'
>>> Color.red = 42
Traceback (most recent call last):
    # ...
Cannot reassign members.
>>> del Color.red
Traceback (most recent call last):
    # ...
AttributeError: Color: cannot delete Enum member.

Note that all attribute changes are disallowed, no new attributes permitted, and deletions are blocked too. When names are enum members, we delegate to the original EnumMeta handling to keep the error messages stable.

If your enum uses properties that alter attributes on the enum class, you'd either have to whitelist those, or allow for names starting with a single underscore to be set; in __setattr__ determine what names would be permissible to set and use super().__setattr__(name, value) for those exceptions, just like the code now distinguishes between class construction and later alterations by using a flag attribute.

The above class can be used just like Enum() to programmatically create an enumeration:

e = FrozenEnum('Things', [('foo',1), ('bar',2)]))

Demo:

>>> e = FrozenEnum('Things', [('foo',1), ('bar',2)])
>>> e
<enum 'Things'>
>>> e.foo = 'bar'
Traceback (most recent call last):
    # ...
AttributeError: Cannot reassign members.

Answer 2

Not necessarily easy, but possible. We need to create a new EnumMeta type¹, create the Enum normally, then reassign the type after the Enum is created:

from enum import Enum, EnumMeta

class FrozenEnum(EnumMeta):
    "prevent creation of new attributes"
    def __getattr__(self, name):
        if name not in self._member_map_:
            raise AttributeError('%s %r has no attribute %r'
                % (self.__class__.__name__, self.__name__, name))
        return super().__getattr__(name)

    def __setattr__(self, name, value):
        if name in self.__dict__ or name in self._member_map_:
            return super().__setattr__(name, value)
        raise AttributeError('%s %r has no attribute %r'
                % (self.__class__.__name__, self.__name__, name))

class Color(Enum):
    red = 1
    green = 2
    blue = 3

Color.__class__ = FrozenEnum

and in use:

>>> type(Color)
<class 'FrozenEnum'>

>>> list(Color)
[<Color.red: 1>, <Color.green: 2>, <Color.blue: 3>]

>>> Color.blue
<Color.blue: 3>

>>> Color.baz = 3
Traceback (most recent call last):
  ...
AttributeError: FrozenEnum 'Color' has no attribute 'baz'

>>> Color.baz
Traceback (most recent call last):
  ...
AttributeError: 'FrozenEnum' object has no attribute 'baz'

Trying to reassign a member still gives the friendlier error:

>>> Color.blue = 9
Traceback (most recent call last):
  ...
AttributeError: Cannot reassign members.

In order to make the class reassignment a little easier, we can write a decorator to encapsulate the process:

def freeze(enum_class):
    enum_class.__class__ = FrozenEnum
    return enum_class

and in use:

@freeze
class Color(Enum):
    red = 1
    green = 2
    blue = 3

Note that it is still possible to overwrite ordinary attributes, such as functions:

@freeze
class Color(Enum):
    red = 1
    green = 2
    blue = 3
    def huh(self):
        print("Huh, I am %s!" % self.name)

and in use:

>>> Color.huh
<function Color.huh at 0x7f7d54ae96a8>

>>> Color.blue.huh()
Huh, I am blue!

>>> Color.huh = 3
>>> Color.huh
3
>>> Color.blue.huh()
Traceback (most recent call last):
  ...
TypeError: 'int' object is not callable

Even that can be blocked, but I'll leave that (for now) as an exercise for someone else.

---

¹ This is only the second case I have seen where subclassing EnumMeta is required. For the other, see this question.

Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.