Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Overload __init__() for a subclass of Enum

I'm trying to overload the __init__() method of a subclass of an enum. Strangely, the pattern that work with a normal class doesn't work anymore with Enum.

The following show the desired pattern working with a normal class:

class Integer:
    def __init__(self, a):
        """Accepts only int"""
        assert isinstance(a, int)
        self.a = a

    def __repr__(self):
        return str(self.a)


class RobustInteger(Integer):
    def __init__(self, a):
        """Accepts int or str"""
        if isinstance(a, str):
            super().__init__(int(a))
        else:
            super().__init__(a)


print(Integer(1))
# 1
print(RobustInteger(1))
# 1
print(RobustInteger('1'))
# 1

The same pattern then breaks if used with an Enum:

from enum import Enum
from datetime import date


class WeekDay(Enum):
    MONDAY = 0
    TUESDAY = 1
    WEDNESDAY = 2
    THURSDAY = 3
    FRIDAY = 4
    SATURDAY = 5
    SUNDAY = 6

    def __init__(self, value):
        """Accepts int or date"""
        if isinstance(value, date):
            super().__init__(date.weekday())
        else:
            super().__init__(value)


assert WeekDay(0) == WeekDay.MONDAY
assert WeekDay(date(2019, 4, 3)) == WeekDay.MONDAY
# ---------------------------------------------------------------------------
# TypeError                                 Traceback (most recent call last)
# /path/to/my/test/file.py in <module>()
#      27 
#      28 
# ---> 29 class WeekDay(Enum):
#      30     MONDAY = 0
#      31     TUESDAY = 1

# /path/to/my/virtualenv/lib/python3.6/enum.py in __new__(metacls, cls, bases, classdict)
#     208             enum_member._name_ = member_name
#     209             enum_member.__objclass__ = enum_class
# --> 210             enum_member.__init__(*args)
#     211             # If another member with the same value was already defined, the
#     212             # new member becomes an alias to the existing one.

# /path/to/my/test/file.py in __init__(self, value)
#      40             super().__init__(date.weekday())
#      41         else:
# ---> 42             super().__init__(value)
#      43 
#      44 

# TypeError: object.__init__() takes no parameters
like image 467
Antoine Gallix Avatar asked Mar 07 '19 19:03

Antoine Gallix


People also ask

What is Auto () in Python?

auto() method, we can get the assigned integer value automatically by just using enum. auto() method. Syntax : enum.auto() Automatically assign the integer value to the values of enum class attributes.

Should enums be capitalized Python?

By convention, enumeration names begin with an uppercase letter and are singular. The enum module is used for creating enumerations in Python.

How do you declare an enum in Python?

Enum is a class in python for creating enumerations, which are a set of symbolic names (members) bound to unique, constant values. The members of an enumeration can be compared by these symbolic anmes, and the enumeration itself can be iterated over. An enum has the following characteristics.

How do you print an enum in Python?

We can access the enum members by using the dot operator(.) with the enum class name. repr() : The repr() method used to print enum member. type(): This is used to print the type of enum member.


2 Answers

You have to overload the _missing_ hook. All instances of WeekDay are created when the class is first defined; WeekDay(date(...)) is an indexing operation rather than a creation operation, and __new__ is initially looking for pre-existing values bound to the integers 0 to 6. Failing that, it calls _missing_, in which you can convert the date object into such an integer.

class WeekDay(Enum):
    MONDAY = 0
    TUESDAY = 1
    WEDNESDAY = 2
    THURSDAY = 3
    FRIDAY = 4
    SATURDAY = 5
    SUNDAY = 6

    @classmethod
    def _missing_(cls, value):
        if isinstance(value, date):
            return cls(value.weekday())
        return super()._missing_(value)

A few examples:

>>> WeekDay(date(2019,3,7))
<WeekDay.THURSDAY: 3>
>>> assert WeekDay(date(2019, 4, 1)) == WeekDay.MONDAY
>>> assert WeekDay(date(2019, 4, 3)) == WeekDay.MONDAY
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AssertionError

(Note: _missing_ is not available prior to Python 3.6.)


Prior to 3.6, it seems you can override EnumMeta.__call__ to make the same check, but I'm not sure if this will have unintended side effects. (Reasoning about __call__ always makes my head spin a little.)

# Silently convert an instance of datatime.date to a day-of-week
# integer for lookup.
class WeekDayMeta(EnumMeta):
    def __call__(cls, value, *args, **kwargs):
        if isinstance(value, date):
            value = value.weekday())
        return super().__call__(value, *args, **kwargs)

class WeekDay(Enum, metaclass=WeekDayMeta):
    MONDAY = 0
    TUESDAY = 1
    WEDNESDAY = 2
    THURSDAY = 3
    FRIDAY = 4
    SATURDAY = 5
    SUNDAY = 6
like image 65
chepner Avatar answered Oct 19 '22 07:10

chepner


There is a much better answer, but I post this anyway as it might be helpful for understanding the issue.

The docs gives this hint:

EnumMeta creates them all while it is creating the Enum class itself, and then puts a custom new() in place to ensure that no new ones are ever instantiated by returning only the existing member instances.

So we have to wait with redefining __new__ until the class is created. With some ugly patching this passes the test:

from enum import Enum
from datetime import date

class WeekDay(Enum):
    MONDAY = 0 
    TUESDAY = 1 
    WEDNESDAY = 2 
    THURSDAY = 3 
    FRIDAY = 4 
    SATURDAY = 5 
    SUNDAY = 6 

wnew = WeekDay.__new__

def _new(cls, value):
    if isinstance(value, date):
        return wnew(cls, value.weekday()) # not date.weekday()
    else:
        return wnew(cls, value)

WeekDay.__new__ = _new

assert WeekDay(0) == WeekDay.MONDAY
assert WeekDay(date(2019, 3, 4)) == WeekDay.MONDAY # not 2019,4,3
like image 37
VPfB Avatar answered Oct 19 '22 08:10

VPfB