bitmask/bitmask.py

#   Copyright 2022 dogeystamp <dogeystamp@disroot.org>
#
#   Redistribution and use in source and binary forms, with or without
#   modification, are permitted provided that the following conditions are met:
#
#   1. Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#
#   2. Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
#   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
#   AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#   IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#   ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
#   LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#   CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
#   SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
#   INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
#   CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
#   ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#   POSSIBILITY OF SUCH DAMAGE.

"""Utilities for manipulating bitmasks."""

def _fullname(obj):
    """Get the full class name of an object, including module.

    Returns:
        String with class name, which can be used in an `eval()`.
    """
    return f"{obj.__class__.__qualname__}"

class Bitmask:
    """Generic bitmask, which represents multiple Enum values.

    Args:
        AllFlags (Enum): Enum of values with corresponding bitmasks.
        flags: Variable length list of flags to enable.

    Examples:
        Initialise bitmask::
            
            from bitmask import Bitmask
            from enum import IntFlag

            class Desc(IntFlag):
                BIG = 1
                ROUND = 1 << 1
                FUNKY = 1 << 2

            bmask = Bitmask(Desc, Desc.BIG, Desc.FUNKY)

        Determine if a flag enabled::

            Desc.FUNKY in bmask

        Show all enabled flags::

            for flag in bmask:
                print(flag)

        Enable flag in bitmask::

            bmask.add(Desc.ROUND)

        Remove flag from bitmask::
            
            bmask.discard(Desc.ROUND)
    """

    def __init__(self, AllFlags, *flags):
        """Init symbol enum and state."""
        self._AllFlags = AllFlags

        self._value = self._AllFlags(0)

        for flag in flags:
            self.add(flag)

    @property
    def AllFlags(self):
        """Enum defining all flags used in the bitmask, and their values."""
        return self._AllFlags

    @property
    def value(self):
        """Integer value for direct access to bitmask state."""
        return self._value

    @value.setter
    def value(self, value):
        if not issubclass(type(value), int):
            raise TypeError(f"value must be an integer (got '{type(value)}')")
        self._value = value
    
    def __contains__(self, item):
        """Determine if flag is enabled."""
        return self.value & item

    def __iter__(self):
        """Return list of enabled flags."""
        for flag in self.AllFlags:
            if flag in self:
                yield flag

    def __int__(self):
        return int(self.value)

    def __str__(self):
        return '|'.join([flag.name for flag in self]) or "0"

    def __repr__(self):
        enum_name = _fullname(self.AllFlags(0))
        args = ', '.join(
            [enum_name] +
            [f"{enum_name}.{flag.name}" for flag in self]
        )

        return f"{_fullname(self)}({args})"

    def _flag_op(self, flag, op):
        """Apply a single flag to the bitmask.

        Args:
            flag (enum value): flag to apply to the bitmask.
            op (function): function that returns new numerical value for the
                bitmask, given the initial value and the flag.
        """
        if not issubclass(type(flag), self.AllFlags):
            raise TypeError(f"can only add {self.AllFlags} (not '{type(flag)}') to {type(self)}")
        self.value = op(self.value, flag)

    def __mask_op(self, other, op):
        """Run operations on two bitmasks/bitmask and flag.

        Will run `__flag_op` with each flag in the other bitmask.

        Args:
            other (Bitmask or Enum value): bitmask to run operation with.
            op (function): operation passed to `__flag_op` for applying individual flags.
        """
        new_bitmask = self.__class__(self.AllFlags, *(flag for flag in self))

        if issubclass(type(other), type(self)):
            for flag in other:
                new_bitmask._flag_op(flag, op)
        elif issubclass(type(other), self.AllFlags):
            new_bitmask._flag_op(flag, op)
        else:
            raise TypeError(f"can only apply {type(self)} or {self.AllFlags} (not '{type(other)}') to {type(self)}")

        return new_bitmask

    def add(self, other):
        """Add single flag to the bitmask."""
        self._flag_op(other, lambda a, b : a | b)

    def __add__(self, other):
        """Union bitmasks/flags together."""
        return self.__mask_op(other, lambda a, b : a | b)

    def __radd__(self, other):
        return self.__add__(other)

    def remove(self, flag):
        """Remove `flag` from the bitmask.

        This behaves the same as built-in `set.remove()`.

        Raises:
            KeyError: flag is not in bitmask.
        """
        if not flag in self:
            raise KeyError(flag)
        self._flag_op(flag, lambda a, b : a ^ b)

    def discard(self, flag):
        """Remove flag bitmask if present.

        This behaves the same as built-in `set.discard()`.
        """
        if flag in self:
            self.remove(flag)