Overriding __contains__ method for a class

Question:

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?

Asked By: clstaudt

||

Answers:

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
Answered By: Devin Jeanpierre

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
Answered By: user1023102
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.