bitarray 3.12.2__cp313-cp313-pyemscripten_2025_0_wasm32.whl

bitarray/util.py

@@ -0,0 +1,694 @@
+# Copyright (c) 2019 - 2026, Ilan Schnell; All Rights Reserved
+# bitarray is published under the PSF license.
+#
+# Author: Ilan Schnell
+"""
+Useful utilities for working with bitarrays.
+"""
+import os
+import sys
+import math
+import random
+import operator
+import itertools
+
+from bitarray import bitarray, frozenbitarray, bits2bytes
+
+from bitarray._util import (
+    zeros, ones, count_n, parity, xor_indices,
+    count_and, count_or, count_xor, any_and, subset,
+    correspond_all, byteswap,
+    serialize, deserialize,
+    ba2hex, hex2ba,
+    ba2base, base2ba,
+    sc_encode, sc_decode,
+    rl_encode, rl_decode,
+    vl_encode, vl_decode,
+    canonical_decode,
+)
+import bitarray._util as _util
+
+__all__ = [
+    'zeros', 'ones', 'urandom', 'random_k', 'random_p', 'gen_primes',
+    'pprint', 'strip', 'count_n',
+    'parity', 'sum_indices', 'xor_indices',
+    'count_and', 'count_or', 'count_xor', 'any_and', 'subset',
+    'correspond_all', 'byteswap', 'intervals',
+    'ba2hex', 'hex2ba',
+    'ba2base', 'base2ba',
+    'ba2int', 'int2ba',
+    'serialize', 'deserialize',
+    'sc_encode', 'sc_decode', 'sc_stat',
+    'rl_encode', 'rl_decode',
+    'vl_encode', 'vl_decode',
+    'huffman_code', 'canonical_huffman', 'canonical_decode',
+]
+
+
+def urandom(__n, endian=None):
+    """urandom(n, /, endian=None) -> bitarray
+
+Return random bitarray of length `n` (uses `os.urandom()`).
+"""
+    a = bitarray(os.urandom(bits2bytes(__n)), endian)
+    del a[__n:]
+    return a
+
+
+def random_k(__n, k, endian=None):
+    """random_k(n, /, k, endian=None) -> bitarray
+
+Return (pseudo-) random bitarray of length `n` with `k` elements
+set to one.  Mathematically equivalent to setting (in a bitarray of
+length `n`) all bits at indices `random.sample(range(n), k)` to one.
+The random bitarrays are reproducible when calling Python's `random.seed()`
+with a specific seed value.
+"""
+    r = _Random(__n, endian)
+    k = operator.index(k)
+    return r.random_k(k)
+
+
+def random_p(__n, p=0.5, endian=None):
+    """random_p(n, /, p=0.5, endian=None) -> bitarray
+
+Return (pseudo-) random bitarray of length `n`, where each bit has
+probability `p` of being one (independent of any other bits).  Mathematically
+equivalent to `bitarray((random() < p for _ in range(n)), endian)`, but much
+faster for large `n`.  The random bitarrays are reproducible when calling
+Python's `random.seed()` with a specific seed value.
+
+This function requires Python 3.12 or higher, as it depends on the standard
+library function `random.binomialvariate()`.  Raises `NotImplementedError`
+when Python version is too low.
+"""
+    if sys.version_info[:2] < (3, 12):
+        raise NotImplementedError("bitarray.util.random_p() requires "
+                                  "Python 3.12 or higher")
+    r = _Random(__n, endian)
+    return r.random_p(p)
+
+
+class _Random:
+
+    # To better understand how the algorithm works, see ./doc/random_p.rst
+    # See also, Internal_Tests and VerificationTests in devel/test_random.py
+
+    # maximal number of calls to .random_half() in .combine_half()
+    M = 8
+
+    # number of resulting probability intervals
+    K = 1 << M
+
+    # limit for setting individual bits randomly
+    SMALL_P = 0.01
+
+    def __init__(self, n=0, endian=None):
+        self.n = n
+        self.nbytes = bits2bytes(n)
+        self.endian = endian
+
+    def random_half(self):
+        """
+        Return bitarray with each bit having probability p = 1/2 of being 1.
+        """
+        nbytes = self.nbytes
+        # Use random module function for reproducibility (not urandom()).
+        # random.randbytes() can be used once Python 3.9 is the minimum
+        # supported version.
+        b = random.getrandbits(8 * nbytes).to_bytes(nbytes, 'little')
+        a = bitarray(b, self.endian)
+        del a[self.n:]
+        return a
+
+    def op_seq(self, i):
+        """
+        Return bitarray containing operator sequence.
+        Each item represents a bitwise operation:   0: AND   1: OR
+        After applying the sequence (see .combine_half()), we
+        obtain a bitarray with probability  q = i / K
+        """
+        if not 0 < i < self.K:
+            raise ValueError("0 < i < %d, got i = %d" % (self.K, i))
+
+        # sequence of &, | operations - least significant operations first
+        a = bitarray(i.to_bytes(2, byteorder="little"), "little")
+        return a[a.index(1) + 1 : self.M]
+
+    def combine_half(self, seq):
+        """
+        Combine random bitarrays with probability 1/2
+        according to given operator sequence.
+        """
+        a = self.random_half()
+        for k in seq:
+            if k:
+                a |= self.random_half()
+            else:
+                a &= self.random_half()
+        return a
+
+    def adjust_count(self, a, increase):
+        """
+        Adjust the number of set bits in bitarray `a` by `increase`.
+        The length of `a` must equal `self.n`, and `increase` must be in the
+        range `-a.count() <= increase <= self.n - a.count()`.  The bitarray
+        is modified in-place at positions selected uniformly at random.
+        """
+        randrange = random.randrange
+        n = self.n
+
+        if increase > 0:  # not enough 1 bits - increase count
+            for _ in range(increase):
+                i = randrange(n)
+                while a[i]:
+                    i = randrange(n)
+                a[i] = 1
+        elif increase < 0:  # too many 1 bits - decrease count
+            for _ in range(-increase):
+                i = randrange(n)
+                while not a[i]:
+                    i = randrange(n)
+                a[i] = 0
+
+    def random_k(self, k):
+        n = self.n
+        # error check inputs and handle edge cases
+        if k <= 0 or k >= n:
+            if k == 0:
+                return zeros(n, self.endian)
+            if k == n:
+                return ones(n, self.endian)
+            raise ValueError("k must be in range 0 <= k <= n, got %s" % k)
+
+        # exploit symmetry to establish: k <= n // 2
+        if k > n // 2:
+            a = self.random_k(n - k)
+            a.invert()  # use in-place to avoid copying
+            return a
+
+        # decide on sequence, see VerificationTests devel/test_random.py
+        if k < 16 or k * self.K < 3 * n:
+            i = 0
+        else:
+            p = k / n  # p <= 0.5
+            p -= (0.2 - 0.4 * p) / math.sqrt(n)
+            i = int(p * (self.K + 1))
+
+        # combine random bitarrays using bitwise AND and OR operations
+        if i < 3:
+            a = zeros(n, self.endian)
+            self.adjust_count(a, k)
+        else:
+            a = self.combine_half(self.op_seq(i))
+            self.adjust_count(a, k - a.count())
+
+        return a
+
+    def random_p(self, p):
+        n = self.n
+        endian = self.endian
+        # error check inputs and handle edge cases
+        if p <= 0.0 or p == 0.5 or p >= 1.0 or not math.isfinite(p):
+            if p == 0.0:
+                return zeros(n, endian)
+            if p == 0.5:
+                return self.random_half()
+            if p == 1.0:
+                return ones(n, endian)
+            raise ValueError("p must be in range 0.0 <= p <= 1.0, got %s" % p)
+
+        # for small n, use literal definition
+        if n < 16:
+            return bitarray((random.random() < p for _ in range(n)), endian)
+
+        # exploit symmetry to establish: p < 0.5
+        if p > 0.5:
+            a = self.random_p(1.0 - p)
+            a.invert()  # use in-place to avoid copying
+            return a
+
+        # for small p, set randomly individual bits
+        if p < self.SMALL_P:
+            a = zeros(n, endian)
+            self.adjust_count(a, random.binomialvariate(n, p))
+            return a
+
+        # calculate operator sequence
+        i = int(p * self.K)
+        if p * (self.K + 1) > i + 1:  # see devel/test_random.py
+            i += 1
+        seq = self.op_seq(i)
+        q = i / self.K
+
+        # when n is small compared to number of operations, also use literal
+        if n < 100 and self.nbytes <= len(seq) + 3 * (q != p):
+            return bitarray((random.random() < p for _ in range(n)), endian)
+
+        # combine random bitarrays using bitwise AND and OR operations
+        a = self.combine_half(seq)
+        if q < p:
+            x = (p - q) / (1.0 - q)
+            a |= self.random_p(x)
+        elif q > p:
+            x = p / q
+            a &= self.random_p(x)
+
+        return a
+
+
+def gen_primes(__n, endian=None, odd=False):
+    """gen_primes(n, /, endian=None, odd=False) -> bitarray
+
+Generate a bitarray of length `n` in which active indices are prime numbers.
+By default (`odd=False`), active indices correspond to prime numbers directly.
+When `odd=True`, only odd prime numbers are represented in the resulting
+bitarray `a`, and `a[i]` corresponds to `2*i+1` being prime or not.
+"""
+    n = operator.index(__n)
+    if n < 0:
+        raise ValueError("bitarray length must be >= 0")
+
+    if odd:
+        a = ones(105, endian)  # 105 = 3 * 5 * 7
+        a[1::3] = 0
+        a[2::5] = 0
+        a[3::7] = 0
+        f = "01110110"
+    else:
+        a = ones(210, endian)  # 210 = 2 * 3 * 5 * 7
+        for i in 2, 3, 5, 7:
+            a[::i] = 0
+        f = "00110101"
+
+    # repeating the array many times is faster than setting the multiples
+    # of the low primes to 0
+    a *= (n + len(a) - 1) // len(a)
+    a[:8] = bitarray(f, endian)
+    del a[n:]
+    # perform sieve starting at 11
+    if odd:
+        for i in a.search(1, 5, int(math.sqrt(n // 2) + 1.0)):  # 11//2 = 5
+            j = 2 * i + 1
+            a[(j * j) // 2 :: j] = 0
+    else:
+        # i*i is always odd, and even bits are already set to 0: use step 2*i
+        for i in a.search(1, 11, int(math.sqrt(n) + 1.0)):
+            a[i * i :: 2 * i] = 0
+    return a
+
+
+def sum_indices(__a, mode=1):
+    """sum_indices(a, /, mode=1) -> int
+
+Return sum of indices of all active bits in bitarray `a`.
+Equivalent to `sum(i for i, v in enumerate(a) if v)`.
+`mode=2` sums square of indices.
+"""
+    # For details see: devel/test_sum_indices.py
+    n = 1 << 19  # block size  512 Kbits
+    if len(__a) <= n:  # shortcut for single block
+        return _util.ssqi(__a, mode)
+
+    if not isinstance(__a, bitarray):
+        raise TypeError("bitarray expected, got '%s'" % type(__a).__name__)
+    mode = operator.index(mode)
+    if mode not in (1, 2):
+        raise ValueError("unexpected mode %r" % mode)
+
+    # Constants
+    m = n // 8  # block size in bytes
+    o1 = n * (n - 1) // 2
+    o2 = o1 * (2 * n - 1) // 3
+
+    nblocks = (len(__a) + n - 1) // n
+    padbits = __a.padbits
+    sm = 0
+    for i in range(nblocks):
+        # use memoryview to avoid copying memory
+        v = memoryview(__a)[i * m : (i + 1) * m]
+        block = bitarray(None, __a.endian, buffer=v)
+        if padbits and i == nblocks - 1:
+            if block.readonly:
+                block = bitarray(block)
+            block[-padbits:] = 0
+
+        k = block.count()
+        if k:
+            y = n * i
+            z1 = o1 if k == n else _util.ssqi(block)
+            if mode == 1:
+                sm += k * y + z1
+            else:
+                z2 = o2 if k == n else _util.ssqi(block, 2)
+                sm += (k * y + 2 * z1) * y + z2
+
+    return sm
+
+
+def pprint(__a, stream=None, group=8, indent=4, width=80):
+    """pprint(bitarray, /, stream=None, group=8, indent=4, width=80)
+
+Pretty-print bitarray object to `stream`, defaults to `sys.stdout`.
+By default, bits are grouped in bytes (8 bits), and 64 bits per line.
+Non-bitarray objects are printed using `pprint.pprint()`.
+"""
+    if stream is None:
+        stream = sys.stdout
+
+    if not isinstance(__a, bitarray):
+        import pprint as _pprint
+        _pprint.pprint(__a, stream=stream, indent=indent, width=width)
+        return
+
+    group = operator.index(group)
+    if group < 1:
+        raise ValueError('group must be >= 1')
+    indent = operator.index(indent)
+    if indent < 0:
+        raise ValueError('indent must be >= 0')
+    width = operator.index(width)
+    if width <= indent:
+        raise ValueError('width must be > %d (indent)' % indent)
+
+    gpl = (width - indent) // (group + 1)  # groups per line
+    epl = group * gpl                      # elements per line
+    if epl == 0:
+        epl = max(1, width - indent - 2)
+    type_name = type(__a).__name__
+    n = len(__a)
+    multiline = n > 0 and (  # here 4 is len("('')")
+        len(type_name) + 4 + n + (n - 1) // group >= width
+    )
+    quotes = "'''" if multiline else ("'" if __a else "")
+
+    stream.write("%s(%s" % (type_name, quotes))
+    for i, b in enumerate(__a):
+        if multiline and i % epl == 0:
+            stream.write('\n%s' % (indent * ' '))
+        if i % group == 0 and i % epl != 0:
+            stream.write(' ')
+        stream.write(str(b))
+
+    if multiline:
+        stream.write('\n')
+
+    stream.write("%s)\n" % quotes)
+    stream.flush()
+
+
+def strip(__a, mode='right'):
+    """strip(bitarray, /, mode='right') -> bitarray
+
+Return a new bitarray with zeros stripped from left, right or both ends.
+Allowed values for mode are the strings: `left`, `right`, `both`
+"""
+    if not isinstance(mode, str):
+        raise TypeError("str expected for mode, got '%s'" %
+                        type(mode).__name__)
+    if mode not in ('left', 'right', 'both'):
+        raise ValueError("mode must be 'left', 'right' or 'both', got %r" %
+                         mode)
+
+    start = None if mode == 'right' else __a.find(1)
+    if start == -1:
+        return __a[:0]
+    stop = None if mode == 'left' else __a.find(1, right=1) + 1
+    return __a[start:stop]
+
+
+def intervals(__a):
+    """intervals(bitarray, /) -> iterator
+
+Compute all uninterrupted intervals of 1s and 0s, and return an
+iterator over tuples `(value, start, stop)`.  The intervals are guaranteed
+to be in order, and their size is always non-zero (`stop - start > 0`).
+"""
+    n = len(__a)
+    if n == 0:
+        return
+
+    value = __a[0]  # value of current interval
+    start = 0
+    while start < n:
+        stop = __a.find(not value, start)
+        if stop < 0:
+            stop = n
+        yield int(value), start, stop
+        start = stop
+        value = not value  # next interval has opposite value
+
+
+def ba2int(__a, signed=False):
+    """ba2int(bitarray, /, signed=False) -> int
+
+Convert the given bitarray to an integer.
+The bit-endianness of the bitarray is respected.
+`signed` indicates whether two's complement is used to represent the integer.
+"""
+    if not isinstance(__a, bitarray):
+        raise TypeError("bitarray expected, got '%s'" % type(__a).__name__)
+
+    length = len(__a)
+    if length == 0:
+        raise ValueError("non-empty bitarray expected")
+
+    res = int.from_bytes(__a.tobytes(), byteorder=__a.endian)
+
+    # Big-endian pad bits are on the right and must be shifted away.
+    if __a.endian == 'big':
+        res >>= __a.padbits
+
+    if signed and res & (1 << (length - 1)):
+        res -= 1 << length
+
+    return res
+
+
+def int2ba(__i, length=None, endian=None, signed=False):
+    """int2ba(int, /, length=None, endian=None, signed=False) -> bitarray
+
+Convert the given integer to a bitarray (with given bit-endianness,
+and no leading (big-endian) / trailing (little-endian) zeros), unless
+the `length` of the bitarray is provided.  An `OverflowError` is raised
+if the integer is not representable with the given number of bits.
+`signed` determines whether two's complement is used to represent the integer,
+and requires `length` to be provided.
+"""
+    i = operator.index(__i)
+    if length is not None:
+        length = operator.index(length)
+        if length <= 0:
+            raise ValueError("length must be > 0")
+
+    if signed:
+        if length is None:
+            raise TypeError("signed requires argument 'length'")
+        m = 1 << (length - 1)
+        if not (-m <= i < m):
+            raise OverflowError("signed integer not in range(%d, %d), "
+                                "got %d" % (-m, m, i))
+        if i < 0:
+            i += 1 << length
+    elif length and i >> length:
+        raise OverflowError("unsigned integer not in range(0, %d), "
+                            "got %d" % (1 << length, i))
+
+    if length is None:
+        length = max(1, i.bit_length())
+
+    a = bitarray(0, endian)
+    a.frombytes(i.to_bytes(bits2bytes(length), byteorder=a.endian))
+
+    if a.endian == 'little':
+        del a[length:]
+    else:
+        del a[:-length]
+    return a
+
+
+def sc_stat(stream):
+    """sc_stat(stream) -> dict
+
+Scan a stream produced by `sc_encode()` and return a dictionary containing
+its bit-endianness (`endian`), length (`nbits`), and the number of blocks
+of each type (`blocks`).  `blocks` is a list such that `blocks[i]` is the
+number of blocks of type `i`.
+
+Like `sc_decode()`, it consumes one encoded bitarray and leaves remaining
+input untouched.
+"""
+    def decode_header(stream):
+        head = next(stream)
+        if head & 0xe0:
+            raise ValueError("invalid header: 0x%02x" % head)
+        endian = 'big' if head & 0x10 else 'little'
+        length = head & 0x0f  # number of bytes representing nbits
+        nbits = 0
+        for j in range(length):
+            nbits |= next(stream) << 8 * j
+        return dict(endian=endian, nbits=nbits)
+
+    def scan_block(stream, blocks):
+        head = next(stream)
+
+        if head < 0xa0:                  # type 0 - 0x00 -- 0x9f
+            if head == 0:  # stop byte
+                return False
+            n = 0
+            k = head if head <= 32 else 32 * (head - 31)
+        elif head < 0xc0:                # type 1 - 0xa0 .. 0xbf
+            n = 1
+            k = head - 0xa0
+        elif 0xc2 <= head <= 0xc4:       # type 2 .. 4 - 0xc2 .. 0xc4
+            n = head - 0xc0
+            k = next(stream)             # index count byte
+        else:
+            raise ValueError("invalid block head: 0x%02x" % head)
+
+        blocks[n] += 1
+
+        nc = max(1, n) * k   # size of block data to consume
+        next(itertools.islice(stream, nc, nc), None)
+        return True
+
+    stream = iter(stream)
+    stats = decode_header(stream)
+
+    blocks = 5 * [0]
+    while scan_block(stream, blocks):
+        pass
+
+    stats['blocks'] = blocks
+    return stats
+
+# ------------------------------ Huffman coding -----------------------------
+
+def _huffman_tree(__freq_map):
+    """_huffman_tree(dict, /) -> Node
+
+Given a dict mapping symbols to their frequency, construct a Huffman tree
+and return its root node.
+"""
+    from heapq import heappush, heappop
+
+    class Node:
+        """
+        There are two types of Node instances (both have 'freq' attribute):
+          * leaf node: has 'symbol' attribute
+          * parent node: has 'child' attribute (tuple with both children)
+        """
+        def __lt__(self, other):
+            # heapq needs to be able to compare the nodes
+            return self.freq < other.freq
+
+    minheap = []
+    # create all leaf nodes and push them onto the queue
+    for sym, f in __freq_map.items():
+        leaf = Node()
+        leaf.symbol = sym
+        leaf.freq = f
+        heappush(minheap, leaf)
+
+    # repeat the process until only one node remains
+    while len(minheap) > 1:
+        # take the two nodes with lowest frequencies from the queue
+        # to construct a new parent node and push it onto the queue
+        parent = Node()
+        parent.child = heappop(minheap), heappop(minheap)
+        parent.freq = parent.child[0].freq + parent.child[1].freq
+        heappush(minheap, parent)
+
+    # the single remaining node is the root of the Huffman tree
+    return minheap[0]
+
+
+def _check_anydict(d):
+    types = dict if sys.version_info[:2] < (3, 15) else (dict, frozendict)
+    if not isinstance(d, types):
+        raise TypeError("dict expected, got '%s'" % type(d).__name__)
+    if not d:
+        raise ValueError("cannot create Huffman code with no symbols")
+
+
+def huffman_code(__freq_map, endian=None):
+    """huffman_code(dict, /, endian=None) -> dict
+
+Given a frequency map, a dictionary mapping symbols to their frequency,
+calculate the Huffman code, i.e. a dict mapping those symbols to
+frozenbitarrays (with given bit-endianness).  Note that the symbols are not
+limited to being strings.  Symbols may be any hashable object.
+"""
+    _check_anydict(__freq_map)
+
+    if len(__freq_map) == 1:
+        # With one symbol, Huffman coding would normally assign it an empty
+        # code.  Here the code must be nonempty, so that .encode() represents
+        # each occurrence and .decode() can recover their number.  We use 0,
+        # leaving 1 unassigned; therefore the prefix code is incomplete.
+        sym = next(iter(__freq_map))
+        return {sym: frozenbitarray('0', endian)}
+
+    result = {}
+
+    def traverse(nd, prefix):
+        try:                    # leaf
+            result[nd.symbol] = prefix
+        except AttributeError:  # parent, so traverse each child
+            traverse(nd.child[0], prefix + '0')
+            traverse(nd.child[1], prefix + '1')
+
+    traverse(_huffman_tree(__freq_map), frozenbitarray(0, endian))
+    return result
+
+
+def canonical_huffman(__freq_map):
+    """canonical_huffman(dict, /) -> tuple
+
+Given a frequency map, a dictionary mapping symbols to their frequency,
+calculate the canonical Huffman code.  Returns a tuple containing:
+
+0. the canonical Huffman code as a dict mapping symbols to frozenbitarrays
+1. a list containing the number of symbols of each code length
+2. a list of symbols in canonical order
+
+Note: the two lists may be used as input for `canonical_decode()`.
+"""
+    _check_anydict(__freq_map)
+
+    if len(__freq_map) == 1:
+        # Only one symbol: see note above in huffman_code()
+        sym = next(iter(__freq_map))
+        return {sym: frozenbitarray('0', 'big')}, [0, 1], [sym]
+
+    code_length = {}  # map symbols to their code length
+
+    def traverse(nd, length):
+        # traverse the Huffman tree, but (unlike in huffman_code() above) we
+        # simply record the length for reaching each symbol
+        try:                    # leaf
+            code_length[nd.symbol] = length
+        except AttributeError:  # parent, so traverse each child
+            traverse(nd.child[0], length + 1)
+            traverse(nd.child[1], length + 1)
+
+    traverse(_huffman_tree(__freq_map), 0)
+
+    # We now have a mapping of symbols to their code length, which is all we
+    # need to construct a list of tuples (symbol, code length) sorted by
+    # code length:
+    table = sorted(code_length.items(), key=lambda item: item[1])
+
+    maxbits = table[-1][1]
+    codedict = {}
+    count = (maxbits + 1) * [0]
+
+    code = 0
+    for i, (sym, length) in enumerate(table):
+        codedict[sym] = frozenbitarray(int2ba(code, length, 'big'))
+        count[length] += 1
+        if i + 1 < len(table):
+            code += 1
+            code <<= table[i + 1][1] - length
+
+    return codedict, count, [item[0] for item in table]