generated from mwc/lab_encoding
122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
from bitarray import bitarray
|
|
from bitarray.util import (
|
|
zeros,
|
|
ones,
|
|
int2ba,
|
|
ba2int,
|
|
)
|
|
from easybits.util import is_bit_string
|
|
from easybits.errors import (
|
|
NotEnoughBits,
|
|
IntegersRequireLength,
|
|
IntegerAdditionRequiresSameLength,
|
|
)
|
|
|
|
class Bits(bitarray):
|
|
"""A wrapper around bitarray with a more idiomatic interface for beginners.
|
|
- Integers are always signed.
|
|
"""
|
|
|
|
default_text_encoding = 'ascii'
|
|
|
|
@classmethod
|
|
def zeros(cls, length):
|
|
return Bits(zeros(length))
|
|
|
|
@classmethod
|
|
def ones(cls, length):
|
|
return Bits(ones(length))
|
|
|
|
def __new__(cls, value=None, length=None, encoding=None):
|
|
if value is None:
|
|
bits = bitarray(length or 0)
|
|
elif isinstance(value, bytes):
|
|
bits = bitarray()
|
|
bits.frombytes(value)
|
|
elif isinstance(value, bool):
|
|
bits = bitarray([value])
|
|
elif isinstance(value, int):
|
|
if not length:
|
|
raise IntegersRequireLength()
|
|
bits = int2ba(value, length=length, signed=True)
|
|
elif isinstance(value, str):
|
|
if encoding:
|
|
bits = bitarray()
|
|
bits.frombytes(value.encode(encoding))
|
|
elif is_bit_string(value):
|
|
bits = bitarray(value)
|
|
else:
|
|
bits = bitarray()
|
|
bits.frombytes(value.encode(cls.default_text_encoding))
|
|
elif isinstance(value, list):
|
|
bits = bitarray(value)
|
|
elif isinstance(value, bitarray):
|
|
bits = value
|
|
else:
|
|
raise ValueError(f"Can't create bits from {value}")
|
|
if length:
|
|
if length < len(bits):
|
|
raise NotEnoughBits(value, length)
|
|
else:
|
|
sized_bits = bitarray(length)
|
|
sized_bits[-len(bits):] = bits
|
|
else:
|
|
sized_bits = bits
|
|
return super().__new__(cls, sized_bits)
|
|
|
|
def __str__(self):
|
|
return self.to01()
|
|
|
|
def __repr__(self):
|
|
return self.to01()
|
|
|
|
@property
|
|
def bool(self):
|
|
return [bool(b) for b in self.tolist()]
|
|
|
|
@property
|
|
def int(self):
|
|
return ba2int(self, signed=True)
|
|
|
|
@property
|
|
def bytes(self):
|
|
return self.tobytes()
|
|
|
|
@property
|
|
def ascii(self):
|
|
return self.bytes.decode("ascii")
|
|
|
|
def __add__(self, other):
|
|
"""Performs bitwise addition on `self` and `other`. Does not
|
|
check for overflow.
|
|
"""
|
|
a, b = self, Bits(other)
|
|
if not len(a) == len(b):
|
|
raise IntegerAdditionRequiresSameLength()
|
|
result = Bits.zeros(len(a))
|
|
carry = 0
|
|
for i in reversed(range(len(a))):
|
|
result[i] = a[i] ^ b[i] ^ carry
|
|
carry = (a[i] & b[i]) | (a[i] & carry) | (b[i] & carry)
|
|
return result
|
|
|
|
def __sub__(self, other):
|
|
"""Performs bitwise subtraction on `self` and `other`. Does not
|
|
check for overflow.
|
|
"""
|
|
a, b = self, Bits(other)
|
|
if not len(a) == len(b):
|
|
raise IntegerAdditionRequiresSameLength()
|
|
return a + (-b)
|
|
|
|
def __neg__(self):
|
|
"""Treats `self` as an integer, and flips its sign.
|
|
"""
|
|
if self[0]:
|
|
return ~self + Bits(1, length=len(self))
|
|
else:
|
|
return ~(self + Bits(-1, length=len(self)))
|
|
|
|
def concat(self, other):
|
|
return super().__add__(other)
|