from __future__ import absolute_import from six import string_types def cmp(a, b): return (a > b) - (a < b) class Bit(object): """ Represents a single Bit. """ def __init__(self, number, is_set=True): self.number = number self.is_set = bool(is_set) self.mask = 2 ** int(number) self.children = [] if not self.is_set: self.mask = ~self.mask def __repr__(self): return "<%s: number=%d, is_set=%s>" % ( self.__class__.__name__, self.number, self.is_set, ) # def __str__(self): # if self.is_set: # return 'Yes' # return 'No' def __int__(self): return self.mask def __bool__(self): return self.is_set __nonzero__ = __bool__ def __eq__(self, value): if isinstance(value, Bit): return value.number == self.number and value.is_set == self.is_set elif isinstance(value, bool): return value == self.is_set elif isinstance(value, int): return value == self.mask return value == self.is_set def __ne__(self, value): return not self == value def __coerce__(self, value): return (self.is_set, bool(value)) def __invert__(self): return self.__class__(self.number, not self.is_set) def __and__(self, value): if isinstance(value, Bit): value = value.mask return value & self.mask def __rand__(self, value): if isinstance(value, Bit): value = value.mask return self.mask & value def __or__(self, value): if isinstance(value, Bit): value = value.mask return value | self.mask def __ror__(self, value): if isinstance(value, Bit): value = value.mask return self.mask | value def __lshift__(self, value): if isinstance(value, Bit): value = value.mask return value << self.mask def __rlshift__(self, value): if isinstance(value, Bit): value = value.mask return self.mask << value def __rshift__(self, value): if isinstance(value, Bit): value = value.mask return value >> self.mask def __rrshift__(self, value): if isinstance(value, Bit): value = value.mask return self.mask >> value def __xor__(self, value): if isinstance(value, Bit): value = value.mask return value ^ self.mask def __rxor__(self, value): if isinstance(value, Bit): value = value.mask return self.mask ^ value def __sentry__(self): return repr(self) def evaluate(self, evaluator, qn, connection): return self.mask, [] def prepare(self, evaluator, query, allow_joins): return evaluator.prepare_node(self, query, allow_joins) class BitHandler(object): """ Represents an array of bits, each as a ``Bit`` object. """ def __init__(self, value, keys, labels=None): # TODO: change to bitarray? if value: self._value = int(value) else: self._value = 0 self._keys = keys self._labels = labels is not None and labels or keys def __eq__(self, other): if not isinstance(other, BitHandler): return False return self._value == other._value def __lt__(self, other): return int(self._value) < other def __le__(self, other): return int(self._value) <= other def __gt__(self, other): return int(self._value) > other def __ge__(self, other): return int(self._value) >= other def __cmp__(self, other): return cmp(self._value, other) def __repr__(self): return "<%s: %s>" % ( self.__class__.__name__, ", ".join( "%s=%s" % (k, self.get_bit(n).is_set) for n, k in enumerate(self._keys) ), ) def __str__(self): return str(self._value) def __int__(self): return self._value def __bool__(self): return bool(self._value) __nonzero__ = __bool__ def __and__(self, value): return BitHandler(self._value & int(value), self._keys) def __or__(self, value): return BitHandler(self._value | int(value), self._keys) def __add__(self, value): return BitHandler(self._value + int(value), self._keys) def __sub__(self, value): return BitHandler(self._value - int(value), self._keys) def __lshift__(self, value): return BitHandler(self._value << int(value), self._keys) def __rshift__(self, value): return BitHandler(self._value >> int(value), self._keys) def __xor__(self, value): return BitHandler(self._value ^ int(value), self._keys) def __contains__(self, key): bit_number = self._keys.index(key) return bool(self.get_bit(bit_number)) def __getattr__(self, key): if key.startswith("_"): return object.__getattribute__(self, key) if key not in self._keys: raise AttributeError("%s is not a valid flag" % key) return self.get_bit(self._keys.index(key)) def __setattr__(self, key, value): if key.startswith("_"): return object.__setattr__(self, key, value) if key not in self._keys: raise AttributeError("%s is not a valid flag" % key) self.set_bit(self._keys.index(key), value) def __iter__(self): return self.iteritems() def __sentry__(self): return repr(self) def _get_mask(self): return self._value mask = property(_get_mask) def evaluate(self, evaluator, qn, connection): return self.mask, [] def get_bit(self, bit_number): mask = 2 ** int(bit_number) return Bit(bit_number, self._value & mask != 0) def set_bit(self, bit_number, true_or_false): mask = 2 ** int(bit_number) if true_or_false: self._value |= mask else: self._value &= ~mask return Bit(bit_number, self._value & mask != 0) def keys(self): return self._keys def iterkeys(self): return iter(self._keys) def items(self): return list(self.iteritems()) def iteritems(self): for k in self._keys: yield (k, getattr(self, k).is_set) def get_label(self, flag): if isinstance(flag, string_types): flag = self._keys.index(flag) if isinstance(flag, Bit): flag = flag.number return self._labels[flag]