Search code examples
pythonpython-dataclasses

How to limit values in dataclass attribute?


I have the following dataclass Gear that I want to limit the maximum value for gear_level from 0 to 5. But as you can see when I increment gear_level, it goes higher than 5, which is not what I want. I tried method as well as postinit. How do I fix this problem?

from dataclasses import dataclass

@dataclass
class Gear:
    gear_level: int = 0
    direction: str = None
    # more codes ...

    def __postinit__(self):
        if self.gear_level <= 0:
            self.gear_level = 0
        elif 5 > self.gear_level > 0:
            self.gear_level = self.gear_level
        else:
            self.gear_level = 5

    def set_gear_level(self, level):
        if level <= 0:
            self.gear_level = 0
        elif 5 > level > 0:
            self.gear_level = level
        else:
            self.gear_level = 5


g = Gear()

g.set_gear_level(6)

print(g)

g.gear_level += 1

print(g)

g.set_gear_level(-1)

print(g)

g.gear_level -= 1
print(g)

Ideally, I prefer to use the g.gear_level += 1 notation, because I want to increment gear_level. It should not jump from gear level 1 to 5. Also, when it decrement, it should stop at 0. It should take both an assignment of 0 and be allowed to decrement to 0. Can this be done?

Gear(gear_level=5, direction=None)
Gear(gear_level=6, direction=None)
Gear(gear_level=0, direction=None)
Gear(gear_level=-1, direction=None)

Solution

  • In this case I would simply use a property:

    @dataclass
    class Gear:
        gear_level: int
    
        # Rest of the class excluded for simplicity
    
        @property
        def gear_level(self) -> int:
            return self._gear_level
    
        @gear_level.setter
        def gear_level(self, value: int) -> None:
            self._gear_level = min(max(value, 0), 5)
    

    This way you don't need to write a __post_init__ or have to remember to call specific methods: assignment to gear_level will be kept 0 <= gear_level <= 5, even with +=.