Search code examples
python-3.xenumstype-safety

Python enum.Enum: Create variables which i can assign enum.Enum members


Creating enumerations in Python 3.4+ is pretty easy:

from enum import Enum

class MyEnum(Enum):
    A = 10
    B = 20

This gets me a typedef MyEnum. With this i can assign a variable:

x = MyEnum.A

So far so good. However things start to get complicate if i like to use enum.Enum's as arguments to functions or class methods and want to assure that class attributes only hold enum.Enum members but not other values.

How can i do this? My idea is sth like this, which i consider more as a workaround than a solution:

class EnContainer:
    def __init__(self, val: type(MyEnum.A) = MyEnum.A):
        assert isinstance(val, type(MyEnum.A))
        self._value = val

Do you have any suggestions or do you see any problems with my approach? I have to consider about 10 different enumerations and would like to come to a consistent approach for initialization, setters and getters.


Solution

  • Instead of type(MyEnum.A), just use MyEnum:

    def __init__(self, val: MyEnum = MyEnum.A):
        assert isinstance(val, MyEnum)
    

    Never use assert for error checking, they are for program validation -- in other words, who is calling EnContainer? If only your own code is calling it with already validated data, then assert is fine; but if code outside your control is calling it, then you should be using proper error checking:

    def __init__(self, val: MyEnum = MyEnum.A):
        if not isinstance(val, MyEnum):
            raise ValueError(
                    "EnContainer called with %s.%r (should be a 'MyEnum')"
                    % (type(val), val)
                    )