16

I need to simulate enums in Python, and did it by writing classes like:

class Spam(Enum):
    k = 3
    EGGS = 0
    HAM = 1
    BAKEDBEANS = 2

Now I'd like to test if some constant is a valid choice for a particular Enum-derived class, with the following syntax: 

if (x in Foo):
    print("seems legit")

Therefore I tried to create an "Enum" base class where I override the __contains__ method like this:

class Enum:
    """
    Simulates an enum.
    """

    k = 0 # overwrite in subclass with number of constants

    @classmethod
    def __contains__(cls, x):
        """
        Test for valid enum constant x:
            x in Enum
        """
        return (x in range(cls.k))

However, when using the in keyword on the class (like the example above), I get the error:

TypeError: argument of type 'type' is not iterable

Why that? Can I somehow get the syntactic sugar I want?

jsbueno
  • 99,910
  • 10
  • 151
  • 209
clstaudt
  • 21,436
  • 45
  • 156
  • 239
  • I think that a clarifiying comment that the code would work fine for: if x in Foo(): print('seems legit') Would make both the question and the answer easier to understand. Because the type(Foo()) is Foo but type(Foo) is the meta class. – phil_20686 Jan 05 '18 at 09:07

2 Answers2

24

Why that?

When you use special syntax like a in Foo, the __contains__ method is looked up on the type of Foo. However, your __contains__ implementation exists on Foo itself, not its type. Foo's type is type, which doesn't implement this (or iteration), thus the error.

The same situation occurs if you instantiate an object and then, after it is created, add a __contains__ function to the instance variables. That function won't be called:

>>> class Empty: pass
... 
>>> x = Empty()
>>> x.__contains__ = lambda: True
>>> 1 in x
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: argument of type 'Empty' is not iterable

Can I somehow get the syntactic sugar I want?

Yes. As mentioned above, the method is looked up on Foo's type. The type of a class is called a metaclass, so you need a new metaclass that implements __contains__.

Try this one:

class MetaEnum(type):
    def __contains__(cls, x):
            return x in range(cls.k)

As you can see, the methods on a metaclass take the metaclass instance -- the class -- as their first argument. This should make sense. It's very similar to a classmethod, except that the method lives on the metaclass and not the class.

Inheritance from a class with a custom metaclass also inherits the metaclass, so you can create a base class like so:

class BaseEnum(metaclass=MetaEnum):
    pass

class MyEnum(BaseEnum):
    k = 3

print(1 in MyEnum) # True
Devin Jeanpierre
  • 92,913
  • 4
  • 55
  • 79
  • 2
    Ah, I had a vague idea that metaclasses exist, but I never had a use case of my own. This works fine, thanks. – clstaudt May 04 '12 at 09:25
  • 8 years later, and this wonderful answer put me on the right track for an issue with I had when switching some packages to Python 3.8. I [posted here](https://stackoverflow.com/a/65225753/758174) what I ended up with, in case it is useful for someone else landing here with a similar problem. – Pierre D Dec 09 '20 at 22:22
1

My usecase was to test on the names of the members of my Enum.

With a slight modification to this solution:

from enum import Enum, EnumMeta, auto


class MetaEnum(EnumMeta):
    def __contains__(cls, item):
        return item in cls.__members__.keys()


class BaseEnum(Enum, metaclass=MetaEnum):
    pass


class LogSections(BaseEnum):
    configuration = auto()
    debug = auto()
    errors = auto()
    component_states = auto()
    alarm = auto()


if __name__ == "__main__":
    print('configuration' in LogSections)
    print('b' in LogSections)

True
False
user1023102
  • 189
  • 3
  • 12