(Python) Treat integer values as a set of flags

class IntFlags:
    """A class for treating integer values as a set of flags"""

    _alias = []

    def __init__(self, state: int = 0):
        if state < 0:
            raise ValueError
        object.__setattr__(self, "_state", state)

    def __getitem__(self, digit: int) -> bool:
        return (self._state & (1 << digit)) > 0

    def __setitem__(self, digit: int, value: bool):
        if type(value) is not bool:
            raise TypeError
        if self[digit] != value:
            object.__setattr__(self, "_state", self._state ^ (1 << digit))

    def __getattr__(self, key: str) -> bool:
        alias = type(self)._alias
        if key not in alias:
            raise AttributeError
        return self[alias.index(key)]

    def __setattr__(self, key: str, value: bool):
        alias = type(self)._alias
        if key not in alias:
            raise AttributeError
        self[alias.index(key)] = value

    def __int__(self) -> int:
        return self._state

    def __str__(self) -> str:
        return f"{self._state:0{len(type(self)._alias)}b}"
    
    def __repr__(self) -> str:
        return f'<{type(self).__name__} "{str(self)}">'

class UnixPermission(IntFlags):
    _alias = ["execute", "write", "read"]
    
    def __str__(self) -> str:
        return "".join(reversed([char if self[idx] else "-" for idx, char in enumerate("xwr")]))

#Initialize
user_permission = UnixPermission(6)
user_permission  #=> <UnixPermission "rw-">

#Access individual flags
user_permission.read  #=> True

#Rewrite individual flags
user_permission.execute = True
user_permission  #=> <UnixPermission "rwx">

#Revert to an integer value
int(user_permission)  #=> 7

Recommended Posts