bitmap.py

class Bitmap:
    def __init__(self, fields, word=None, debug=False):
        """
        Create a bitmap for the specified list of field names.
        Each field is a bit in the associated map.  The list is specified
        from least to most significant bit and all field names must
        be unique.

        :param fields:  list of field names from least significant bit to
           most significant bit
        :param word:  initial value for word
        :param debug:  If True, prints debugging messages
        """

        self.fields = fields
        self.bitsN = len(fields)
        self.masks = [2**b for b in range(self.bitsN)]
        self.name2mask = dict(zip(fields, self.masks))
        self.name2bit = dict(zip(fields, range(self.bitsN)))
        self.word = word
        self.debug = debug

        if self.debug:
            print(f"{self.bitsN} bit word, initial value {bin(self.word)}")

    def _complement(self, val):
        """
        Python only has one int type.
        Bit-wise negation performs complement as if the word size was
        one bit longer than the most significant bit that is set.
        This complement acts as if the word size is based on the
        number of fields in the class.

        :param val:  value to be complemented
        :return:  ones complement of value
        """
        # Set one past highest order bit
        high = val | 2**self.bitsN
        comp = ~high   # complement
        return comp & (2**self.bitsN - 1)  # return bits in range

    def __str__(self):
        """
        Return field names that are set
        Ordered from most significant bit to least significant bit
        """
        val = self.word
        idx = 0
        fields = []
        while val > 0:
            bit_set = val & 1
            if bit_set:
                # store the set field coerced to a string
                fields.append(str(self.fields[idx]))
            # next field and right shift
            idx = idx + 1
            val = val >> 1

        fields.reverse()  # reorder most significant bit first
        return ", ".join(fields)

    def set_word(self, word):
        """
        Override the bitmask value with specified word
        :param word:  new value
        :return: None
        """
        self.word = word

        if self.debug:
            print(f"Setting bitfield to {bin(self.word)}")

    def is_set(self, field):
        """
        Check if field is set
        :param field:
        :return: bool
        """
        if self.word is None:
            raise ValueError("Must create with word or call set_word first")
        try:
            value = self.name2mask[field] & self.word > 0
        except KeyError as k:
            field_names = ', '.join(self.fields)
            raise ValueError(
                f"Unknown field: '{field}', expected: one of: {field_names}"
            ) from k
        return self.name2mask[field] & self.word > 0

    def set(self, field, value):
        """
        Set field to value
        :param field:  bitfiled name
        :param value:  True|non zero-> set, False|zero -> clear
        """

        if self.word is None:
            raise ValueError("Must create with word or call set_word first")
        mask = self.name2mask[field]
        if value:
            self.word = self.word | mask  # set bit
        else:
            self.word = self.word & self._complement(mask)  # clear bit

        if self.debug:
            print(f"set bit {self.name2bit[field]} to {value}: {bin(self.word)}")