I have an Enum class called Product like this:

class Product(Enum):
   HELMET = 1
   VEHICLE = 2
   CAR = 3
   TRUCK = 4

And i do some comparisons between two Products. As you can see, however, some items are "subcategories" of others (CAR and TRUCK can also be considered VEHICLE). In order to prevent bugs in my code, i came up with overriding __eq__ method:

class Product(Enum):
   CLOTH = 1
   VEHICLE = 2
   CAR = 3
   TRUCK = 4

   def __eq__(self, other):
      
      if not isinstance(other, Product):
         return False
      
      if Product.CLOTH not in (self.name, other.name):
         if Product.VEHICLE in (self.name, other.name):
            return True
         return False
      return self.value == other.value

• First, i compare if other is an instance of Product, returning false if it's not.

• Then, i check if one (or both) of the items to be compared are CLOTH.

  • If not, i check if one (or both) are VEHICLE.
    • If is, then i'm comparing VEHICLE with TRUCK or CAR, which is True.
    • If not, then i'm comparing TRUCK with CAR, which is False.

• If code reaches the final return self.value == other.value, it means that i'm comparing CLOTH with something else, so it's a basic check based in value of the Enums.

I'd like to know if there's a better approach for this, or a way that i can optimize those if statements.

Test outlines

Product.CLOTH == Product.VEHICLE # must return False
Product.CLOTH == Product.TRUCK # must return False
Product.CLOTH == Product.CAR # must return False

Product.CAR == Product.VEHICLE # must return True
Product.CAR == Product.TRUCK # must return False

Thanks in advance.

CodePudding user response：

you can try a thing like this:

def __eq__(self, other):
    if not isinstance(other, Product):
        return False

    if Product.CLOTH in (self.name, other.name):
        return self.value == other.value

    return Product.VEHICLE in (self.name, other.name)

CodePudding user response：

It would make more sense to use classes here I think

class Vehicle(object):
    pass

class Cloth(object):
    pass

class Car(Vehicle):
    pass

class Truck(Vehicle):
    pass

class Helmet(Cloth):
    pass

car = Car()
truck = Truck()
helmet = Helmet()

print(f"comparing car, Vehicle: {isinstance(car, Vehicle)}")
print(f"comparing truck, Vehicle: {isinstance(truck, Vehicle)}")
print(f"comparing truck, Car: {isinstance(truck, Car)}")
print(f"comparing helmet, Cloth: {isinstance(helmet, Cloth)}")
print(f"comparing helmet, Car: {isinstance(helmet, Car)}")

gives

comparing car, Vehicle: True
comparing truck, Vehicle: True
comparing truck, Car: False
comparing helmet, Cloth: True
comparing helmet, Car: False