Search code examples
pythonenumstype-safety

How to isolate enums by making a distinction on type?


The following code defines two enums

class Insect:
    BEE = 0x00
    WASP = 0x01
    BUMBLEBEE = 0x02


class Breakfast:
    HAM = 0x00
    EGGS = 0x01
    PANCAKES = 0x02


b = Insect.WASP
if b == Breakfast.EGGS:
    print("ok")

As the condition illustrates, one is left open to make a mistake of testing against an entirely distinct enum. How can I isolate the enums, by type not by distinct values, so that the test above will generate an error?

Update

I see that this is one of the finer points along the road of moving from Python 2 to Python 3.

Thanks to wim's suggestion, the following code will generate an error if I attempt to compare apples and oranges.

from enum import Enum


class Apple(Enum):
    RED_DELICIOUS = 0x00
    GALA = 0x01
    FUJI = 0x02

    def __eq__(self, other):
        if type(other) is not type(self):
            raise Exception("You can't compare apples and oranges.")
        return super().__eq__(other)


class Orange(Enum):
    NAVEL = 0x00
    BLOOD = 0x01
    VALENCIA = 0x02

    def __eq__(self, other):
        if type(other) is not type(self):
            raise Exception("You can't compare apples and oranges.")
        return super().__eq__(other)


apple = Apple.GALA
if apple == Orange.BLOOD:
    print("ok")

Solution

  • Don't use a custom class for it. Use the stdlib's enum type, they'll do the right thing here.

    from enum import Enum
    
    class Insect(Enum):
        ...
    

    If you want a hard crash:

    class MyEnum(Enum):
    
        def __eq__(self, other):
            if type(other) is not type(self):
                raise Exception("Don't do that")
            return super().__eq__(other)
    

    But I caution against this design, since:

    1. enum instances are often compared by identity and not equality
    2. there is little (no?) precedent in Python for equality comparisons raising errors