BitMask

bitmask.py

"""BitMask"""

class BitMask:
    __slots__ = ("size", "mask")
    
    def __init__(self, size: int = 16):
        """Create a bit mask to store (size) 0/1 status"""
        self.size = size
        self.mask = 1 << size
        
    def __repr__(self) -> str:
        return f"BitMask(size={self.size}, mask={self.mask})"
        
    def __hash__(self) -> int:
        return hash(self.mask)

    def switch_on(self, i: int):
        """Set bit status on position i to 1"""
        assert i < self.size, "i overflow"
        self.mask |= 1 << i

    def switch_off(self, i: int):
        """Set bit status on position i to 0"""
        if self[i]:
            self.mask ^= 1 << i

    def __getitem__(self, i: int) -> bool:
        """Return the bit status on ith position"""
        return bool(1 & (self.mask >> i))