如何简化Python类初始化方法中断言语句的多个or判断逻辑？

Question #

I've written the __init__ method for a Python Weapon class, using multiple or operators to assert parameter validity:

class weapon:
    def __init__(self, name: str, weapontype : str, range : str, material : str, price : float, quantity : int):
        assert weapontype == "bladed" or weapontype == "blunt" or weapontype == "polearm" or weapontype == "Ranged", f"{weapontype} is not a valid weapontype"
        assert range == "close range" or range == "medium range" or range == "long range", f"{range} is not a valid range"
        assert material == "metal" or material == "wood", f"{material} is not a valid material type"
        self.name = name
        self.weapontype = weapontype
        self.range = range
        self.material = material
        self.price = price
        self.quantity = quantity

Is there a cleaner way to replace these or operators? I've looked up relevant information but haven't found a solution.

Answer #

Absolutely! Those chained or checks can get messy and hard to maintain as your valid options grow—here are two cleaner, more scalable approaches:

1. Use in with a set of valid values #

This is the simplest fix: define your valid options as a set (for O(1) lookup speed) and check if the input value is in that set. It’s concise, easy to update, and you can even improve error messages to list valid options for better debugging.

class Weapon:
    def __init__(self, name: str, weapontype: str, range: str, material: str, price: float, quantity: int):
        # Define valid options as sets (faster lookups than lists)
        VALID_WEAPON_TYPES = {"bladed", "blunt", "polearm", "Ranged"}
        VALID_RANGES = {"close range", "medium range", "long range"}
        VALID_MATERIALS = {"metal", "wood"}
        
        assert weapontype in VALID_WEAPON_TYPES, f"{weapontype} is not a valid weapontype. Valid options: {', '.join(VALID_WEAPON_TYPES)}"
        assert range in VALID_RANGES, f"{range} is not a valid range. Valid options: {', '.join(VALID_RANGES)}"
        assert material in VALID_MATERIALS, f"{material} is not a valid material type. Valid options: {', '.join(VALID_MATERIALS)}"
        
        self.name = name
        self.weapontype = weapontype
        self.range = range
        self.material = material
        self.price = price
        self.quantity = quantity

2. Use enum.Enum for type safety #

If you want stronger type enforcement (to avoid accidental typos like "blade" instead of "bladed"), using Python’s Enum class is a great choice. It encapsulates valid values, gives you IDE autocompletion, and makes your code more self-documenting.

from enum import Enum

# Define enums for each category of valid values
class WeaponType(Enum):
    BLADED = "bladed"
    BLUNT = "blunt"
    POLEARM = "polearm"
    RANGED = "Ranged"

class WeaponRange(Enum):
    CLOSE = "close range"
    MEDIUM = "medium range"
    LONG = "long range"

class Material(Enum):
    METAL = "metal"
    WOOD = "wood"

class Weapon:
    def __init__(self, name: str, weapontype: WeaponType, range: WeaponRange, material: Material, price: float, quantity: int):
        # Optional: Handle cases where someone might pass a string instead of the Enum
        if isinstance(weapontype, str):
            weapontype = WeaponType(weapontype)
        if isinstance(range, str):
            range = WeaponRange(range)
        if isinstance(material, str):
            material = Material(material)
        
        # Store either the Enum instance or its string value (your choice)
        self.name = name
        self.weapontype = weapontype.value
        self.range = range.value
        self.material = material.value
        self.price = price
        self.quantity = quantity

With enums, you get runtime checks that the input is valid, and you don’t need explicit assertions if you enforce the type hints (though the string fallback is helpful for flexibility).

Which to choose? #

• Use the set + in approach for quick, simple validation where you don’t need strict type enforcement.
• Use Enums if you want to enforce valid values throughout your codebase, reduce typos, and make your code more readable.

内容的提问来源于stack exchange，提问作者user19803102