bitarray 3.12.2__cp313-cp313-pyemscripten_2025_0_wasm32.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
bitarray/test_util.py ADDED
@@ -0,0 +1,3166 @@
1
+ # Copyright (c) 2019 - 2026, Ilan Schnell; All Rights Reserved
2
+ # bitarray is published under the PSF license.
3
+ #
4
+ # Author: Ilan Schnell
5
+ """
6
+ Tests for bitarray.util module
7
+ """
8
+ import os
9
+ import sys
10
+ import math
11
+ import array
12
+ import base64
13
+ import binascii
14
+ import operator
15
+ import struct
16
+ import tempfile
17
+ import unittest
18
+ from io import StringIO
19
+ from functools import reduce
20
+ from random import (choice, choices, getrandbits, randrange, randint, random,
21
+ sample, seed)
22
+ from string import hexdigits, whitespace
23
+ from collections import Counter
24
+
25
+ from bitarray import (bitarray, frozenbitarray, decodetree, bits2bytes,
26
+ get_default_endian)
27
+ from bitarray.test_bitarray import Util, is_pypy, urandom_2, ENDIANS, PTRSIZE
28
+
29
+ from bitarray.util import (
30
+ zeros, ones, urandom, random_k, random_p, pprint, strip, count_n,
31
+ parity, gen_primes, sum_indices, xor_indices,
32
+ count_and, count_or, count_xor, any_and, subset,
33
+ correspond_all, byteswap, intervals,
34
+ serialize, deserialize, ba2hex, hex2ba, ba2base, base2ba,
35
+ ba2int, int2ba,
36
+ sc_encode, sc_decode, sc_stat, rl_encode, rl_decode, vl_encode, vl_decode,
37
+ huffman_code, canonical_huffman, canonical_decode,
38
+ )
39
+ from bitarray._util import uleb128_encode # type: ignore
40
+
41
+
42
+ OPT_ENDIANS = ENDIANS + (None,)
43
+
44
+ # --------------------------- zeros() ones() -----------------------------
45
+
46
+ class ZerosOnesTests(unittest.TestCase):
47
+
48
+ def test_basic(self):
49
+ for _ in range(50):
50
+ a = choice([zeros(0), zeros(0, None), zeros(0, endian=None),
51
+ ones(0), ones(0, None), ones(0, endian=None)])
52
+ self.assertEqual(a, bitarray())
53
+ self.assertEqual(a.endian, get_default_endian())
54
+ self.assertIs(type(a), bitarray)
55
+
56
+ endian = choice(OPT_ENDIANS)
57
+ n = randrange(100)
58
+
59
+ a = choice([zeros(n, endian), zeros(n, endian=endian)])
60
+ self.assertEqual(a.to01(), n * "0")
61
+ self.assertEqual(a.endian, endian or get_default_endian())
62
+
63
+ b = choice([ones(n, endian), ones(n, endian=endian)])
64
+ self.assertEqual(b.to01(), n * "1")
65
+ self.assertEqual(b.endian, endian or get_default_endian())
66
+
67
+ def test_errors(self):
68
+ for f in zeros, ones:
69
+ self.assertRaises(TypeError, f) # no argument
70
+ self.assertRaises(TypeError, f, '')
71
+ self.assertRaises(TypeError, f, bitarray())
72
+ self.assertRaises(TypeError, f, [])
73
+ self.assertRaises(TypeError, f, 1.0)
74
+ self.assertRaises(ValueError, f, -1)
75
+
76
+ # endian not string
77
+ for x in 0, 1, {}, [], False, True:
78
+ self.assertRaises(TypeError, f, 0, x)
79
+ # endian wrong string
80
+ self.assertRaises(ValueError, f, 0, 'foo')
81
+
82
+ # ----------------------------- urandom() ---------------------------------
83
+
84
+ class URandomTests(unittest.TestCase):
85
+
86
+ def test_basic(self):
87
+ for _ in range(20):
88
+ a = choice([urandom(0), urandom(0, endian=None)])
89
+ self.assertEqual(a, bitarray())
90
+ self.assertEqual(a.endian, get_default_endian())
91
+
92
+ endian = choice(OPT_ENDIANS)
93
+ n = randrange(100)
94
+
95
+ a = choice([urandom(n, endian), urandom(n, endian=endian)])
96
+ self.assertEqual(len(a), n)
97
+ self.assertEqual(a.endian, endian or get_default_endian())
98
+ self.assertIs(type(a), bitarray)
99
+
100
+ def test_errors(self):
101
+ U = urandom
102
+ self.assertRaises(TypeError, U)
103
+ self.assertRaises(TypeError, U, '')
104
+ self.assertRaises(TypeError, U, bitarray())
105
+ self.assertRaises(TypeError, U, [])
106
+ self.assertRaises(TypeError, U, 1.0)
107
+ self.assertRaises(ValueError, U, -1)
108
+ self.assertRaises(TypeError, U, 0, 1)
109
+ self.assertRaises(ValueError, U, 0, 'foo')
110
+
111
+ def test_count(self):
112
+ a = urandom(10_000_000)
113
+ # see if population is within expectation
114
+ self.assertTrue(abs(a.count() - 5_000_000) <= 15_811)
115
+
116
+ # ---------------------------- random_k() ---------------------------------
117
+
118
+ class Random_K_Tests(unittest.TestCase):
119
+
120
+ def test_basic(self):
121
+ for _ in range(250):
122
+ endian = choice(OPT_ENDIANS)
123
+ n = randrange(120)
124
+ k = randint(0, n)
125
+ a = random_k(n, k, endian)
126
+ self.assertIs(type(a), bitarray)
127
+ self.assertEqual(len(a), n)
128
+ self.assertEqual(a.count(), k)
129
+ self.assertEqual(a.endian, endian or get_default_endian())
130
+
131
+ def test_inputs_and_edge_cases(self):
132
+ R = random_k
133
+ self.assertRaises(TypeError, R)
134
+ self.assertRaises(TypeError, R, 4)
135
+ self.assertRaises(TypeError, R, 1, "0.5")
136
+ self.assertRaises(TypeError, R, 10, 4.0)
137
+ self.assertRaises(TypeError, R, 10.0, 4)
138
+ self.assertRaises(TypeError, R, 1, p=1)
139
+ self.assertRaises(TypeError, R, 11, 5.5) # see issue #239
140
+ self.assertRaises(ValueError, R, -1, 0)
141
+ for k in -1, 11: # k is not 0 <= k <= n
142
+ self.assertRaises(ValueError, R, 10, k)
143
+ self.assertRaises(ValueError, R, 10, 7, 'foo')
144
+ self.assertRaises(ValueError, R, 10, 7, endian='foo')
145
+ for n in range(20):
146
+ self.assertEqual(R(n, k=0), zeros(n))
147
+ self.assertEqual(R(n, k=n), ones(n))
148
+
149
+ def test_count(self):
150
+ for n in range(10): # test explicitly for small n
151
+ for k in range(n + 1):
152
+ a = random_k(n, k)
153
+ self.assertEqual(len(a), n)
154
+ self.assertEqual(a.count(), k)
155
+
156
+ for _ in range(100):
157
+ n = randrange(10_000)
158
+ k = randint(0, n)
159
+ a = random_k(n, k)
160
+ self.assertEqual(len(a), n)
161
+ self.assertEqual(a.count(), k)
162
+
163
+ def test_active_bits(self):
164
+ # test if all bits are active
165
+ n = 240
166
+ cum = zeros(n)
167
+ for _ in range(1000):
168
+ k = randint(30, 40)
169
+ a = random_k(n, k)
170
+ self.assertEqual(a.count(), k)
171
+ cum |= a
172
+ if cum.all():
173
+ break
174
+ else:
175
+ self.fail()
176
+
177
+ @unittest.skipIf(sys.version_info[:2] < (3, 8), "math.comb added in 3.8")
178
+ def test_combinations(self):
179
+ # for entire range of 0 <= k <= n, validate that random_k()
180
+ # generates all possible combinations
181
+ n = 7
182
+ total = 0
183
+ for k in range(n + 1):
184
+ expected = math.comb(n, k)
185
+ combs = set()
186
+ for _ in range(10_000):
187
+ combs.add(frozenbitarray(random_k(n, k)))
188
+ if len(combs) == expected:
189
+ total += expected
190
+ break
191
+ else:
192
+ self.fail()
193
+ self.assertEqual(total, 2 ** n)
194
+
195
+ def collect_code_branches(self):
196
+ # return list of bitarrays from all code branches of random_k()
197
+ res = []
198
+ # test small k (no .combine_half())
199
+ res.append(random_k(300, 10))
200
+ # general cases
201
+ for k in 100, 500, 2_500, 4_000:
202
+ res.append(random_k(5_000, k))
203
+ return res
204
+
205
+ def test_seed(self):
206
+ # We ensure that after setting a seed value, random_k() will
207
+ # always return the same random bitarrays. However, we do not ensure
208
+ # that these results will not change in future versions of bitarray.
209
+ a = []
210
+ for val in 654321, 654322, 654321, 654322:
211
+ seed(val)
212
+ a.append(self.collect_code_branches())
213
+ self.assertEqual(a[0], a[2])
214
+ self.assertEqual(a[1], a[3])
215
+ for item0, item1 in zip(a[0], a[1]):
216
+ self.assertNotEqual(item0, item1)
217
+ # initialize seed with current system time again
218
+ seed()
219
+
220
+
221
+ # ---------------------------- random_p() ---------------------------------
222
+
223
+ HAVE_BINOMIALVARIATE = sys.version_info[:2] >= (3, 12)
224
+
225
+ @unittest.skipIf(HAVE_BINOMIALVARIATE, "test for Python <3.12 - "
226
+ "random.binomialvariate() missing")
227
+ class Random_P_Not_Implemented(unittest.TestCase):
228
+
229
+ def test_not_implemented(self):
230
+ self.assertRaises(NotImplementedError, random_p, 100, 0.25)
231
+
232
+
233
+ @unittest.skipIf(not HAVE_BINOMIALVARIATE, "Python 3.12+ required")
234
+ class Random_P_Tests(unittest.TestCase):
235
+
236
+ def test_basic(self):
237
+ for _ in range(250):
238
+ endian = choice(OPT_ENDIANS)
239
+ n = randrange(120)
240
+ p = choice([0.0, 0.0001, 0.2, 0.5, 0.9, 1.0])
241
+ a = random_p(n, p, endian)
242
+ self.assertIs(type(a), bitarray)
243
+ self.assertEqual(len(a), n)
244
+ self.assertEqual(a.endian, endian or get_default_endian())
245
+
246
+ def test_inputs_and_edge_cases(self):
247
+ R = random_p
248
+ self.assertRaises(TypeError, R)
249
+ self.assertRaises(TypeError, R, 0.25)
250
+ self.assertRaises(TypeError, R, 1, "0.5")
251
+ self.assertRaises(ValueError, R, -1)
252
+ self.assertRaises(ValueError, R, 1, -0.5)
253
+ self.assertRaises(ValueError, R, 1, p=1.5)
254
+ self.assertRaises(ValueError, R, 1, 0.15, 'foo')
255
+ self.assertRaises(ValueError, R, 10, 0.5, endian='foo')
256
+ self.assertEqual(R(0), bitarray())
257
+ for n in range(30):
258
+ self.assertEqual(R(n, 0), zeros(n))
259
+ self.assertEqual(len(R(n, 0.5)), n)
260
+ self.assertEqual(R(n, p=1), ones(n))
261
+ self.assertRaises(ValueError, R, n, math.nan)
262
+
263
+ def test_default(self):
264
+ a = random_p(10_000_000) # p defaults to 0.5
265
+ # see if population is within expectation
266
+ self.assertTrue(abs(a.count() - 5_000_000) <= 15_811)
267
+
268
+ def test_count(self):
269
+ for _ in range(500):
270
+ n = choice([randrange(4, 120), randrange(100, 1000)])
271
+ p = choice([0.0001, 0.001, 0.01, 0.1, 0.25, 0.5, 0.9])
272
+ sigma = math.sqrt(n * p * (1.0 - p))
273
+ a = random_p(n, p)
274
+ self.assertEqual(len(a), n)
275
+ self.assertTrue(abs(a.count() - n * p) < max(4, 10 * sigma))
276
+
277
+ def collect_code_branches(self):
278
+ # return list of bitarrays from all code branches of random_p()
279
+ res = []
280
+ # for default p=0.5, random_p uses getrandbits
281
+ res.append(random_p(32))
282
+ # test small p
283
+ res.append(random_p(5_000, 0.002))
284
+ # small n (note that p=0.4 will call the "literal definition" case)
285
+ res.append(random_p(15, 0.4))
286
+ # general cases
287
+ for p in 0.1, 0.2, 0.375, 0.4999, 0.7:
288
+ res.append(random_p(150, p))
289
+ return res
290
+
291
+ def test_seed(self):
292
+ # We ensure that after setting a seed value, random_p() will always
293
+ # return the same random bitarrays. However, we do not ensure that
294
+ # these results will not change in future versions of bitarray.
295
+ a = []
296
+ for val in 123456, 123457, 123456, 123457:
297
+ seed(val)
298
+ a.append(self.collect_code_branches())
299
+ self.assertEqual(a[0], a[2])
300
+ self.assertEqual(a[1], a[3])
301
+ for item0, item1 in zip(a[0], a[1]):
302
+ self.assertNotEqual(item0, item1)
303
+ # initialize seed with current system time again
304
+ seed()
305
+
306
+
307
+ # ---------------------------- gen_primes() -------------------------------
308
+
309
+ class PrimeTests(unittest.TestCase):
310
+
311
+ primes = [
312
+ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61,
313
+ 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137,
314
+ 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211,
315
+ 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271, 277, 281, 283,
316
+ 293, 307, 311, 313, 317, 331, 337, 347, 349, 353, 359, 367, 373, 379,
317
+ 383, 389, 397, 401, 409, 419, 421, 431, 433, 439, 443, 449, 457, 461,
318
+ ]
319
+
320
+ def test_errors(self):
321
+ P = gen_primes
322
+ self.assertRaises(TypeError, P, 3, 1)
323
+ self.assertRaises(TypeError, P, "1")
324
+ self.assertRaises(TypeError, P, "1.0")
325
+ self.assertRaises(TypeError, P, 123.0)
326
+ self.assertRaises(ValueError, P, -1)
327
+ self.assertRaises(TypeError, P, 8, 4)
328
+ self.assertRaises(TypeError, P, 8, foo="big")
329
+ self.assertRaises(ValueError, P, 8, "foo")
330
+ self.assertRaises(ValueError, P, 8, endian="foo")
331
+
332
+ def test_explicit(self):
333
+ for n in range(230):
334
+ endian = choice(OPT_ENDIANS)
335
+ odd = getrandbits(1)
336
+ a = gen_primes(n, endian, odd)
337
+ self.assertEqual(len(a), n)
338
+ self.assertEqual(a.endian, endian or get_default_endian())
339
+ if odd:
340
+ lst = [2] + [2 * i + 1 for i in a.search(1)]
341
+ else:
342
+ lst = [i for i in a.search(1)]
343
+ self.assertEqual(lst, self.primes[:len(lst)])
344
+
345
+ def test_cmp(self):
346
+ N = 10_000
347
+ c = ones(N)
348
+ c[:2] = 0
349
+ for i in range(int(math.sqrt(N) + 1.0)):
350
+ if c[i]:
351
+ c[i * i :: i] = 0
352
+ self.assertEqual(list(c.search(1, 0, 462)), self.primes)
353
+
354
+ for _ in range(20):
355
+ n = randrange(N)
356
+ endian = choice(ENDIANS)
357
+ a = gen_primes(n, endian=endian)
358
+ self.assertEqual(a, c[:n])
359
+ self.assertEqual(a.endian, endian)
360
+
361
+ b = gen_primes(n // 2, endian, odd=True)
362
+ self.assertEqual(b, a[1::2])
363
+ self.assertEqual(b, c[1:n:2])
364
+
365
+ for _ in range(20):
366
+ i = randrange(10, 100)
367
+ x = randint(-1, 1)
368
+ n = i * i + x
369
+ self.assertEqual(gen_primes(n), c[:n])
370
+ self.assertEqual(gen_primes(n // 2, odd=1), c[1:n:2])
371
+
372
+ self.assertEqual(gen_primes(N), c)
373
+ self.assertEqual(gen_primes(N // 2, odd=1), c[1::2])
374
+
375
+ def test_count(self):
376
+ for n, count, sum_p, sum_sqr_p in [
377
+ ( 10, 4, 17, 87),
378
+ ( 100, 25, 1_060, 65_796),
379
+ ( 1_000, 168, 76_127, 49_345_379),
380
+ (10_000, 1229, 5_736_396, 37_546_387_960),
381
+ ]:
382
+ a = gen_primes(n)
383
+ self.assertEqual(len(a), n)
384
+ self.assertEqual(a.count(), count)
385
+ self.assertEqual(sum_indices(a), sum_p)
386
+ self.assertEqual(sum_indices(a, 2), sum_sqr_p)
387
+ b = gen_primes(n // 2, odd=1)
388
+ self.assertEqual(len(b), n // 2)
389
+ self.assertEqual(b.count() + 1, count) # +1 because of prime 2
390
+ self.assertEqual(b, a[1::2])
391
+
392
+ # ----------------------------- pprint() ----------------------------------
393
+
394
+ class PPrintTests(unittest.TestCase):
395
+
396
+ def test_example(self):
397
+ a = gen_primes(100)
398
+ f = StringIO()
399
+ pprint(a, f, indent=3, group=10, width=60)
400
+ self.assertEqual(
401
+ f.getvalue(), "bitarray('''\n"
402
+ " 0011010100 0101000101 0001000001 0100000100 0101000100\n"
403
+ " 0001000001 0100000100 0101000001 0001000001 0000000100\n"
404
+ "''')\n")
405
+
406
+ @staticmethod
407
+ def get_code_string(a):
408
+ f = StringIO()
409
+ pprint(a, stream=f)
410
+ return f.getvalue()
411
+
412
+ def round_trip(self, a):
413
+ b = eval(self.get_code_string(a))
414
+ self.assertEqual(b, a)
415
+ self.assertIs(type(b), type(a))
416
+
417
+ def test_bitarray(self):
418
+ a = bitarray('110')
419
+ self.assertEqual(self.get_code_string(a), "bitarray('110')\n")
420
+ self.round_trip(a)
421
+
422
+ def test_frozenbitarray(self):
423
+ a = frozenbitarray('01')
424
+ self.assertEqual(self.get_code_string(a), "frozenbitarray('01')\n")
425
+ self.round_trip(a)
426
+
427
+ def test_formatting(self):
428
+ a = bitarray(200)
429
+ for width in range(40, 130, 10):
430
+ for n in range(1, 10):
431
+ f = StringIO()
432
+ pprint(a, stream=f, group=n, width=width)
433
+ r = f.getvalue()
434
+ self.assertEqual(eval(r), a)
435
+ s = r.strip("bitary(')\n")
436
+ for group in s.split()[:-1]:
437
+ self.assertEqual(len(group), n)
438
+ for line in s.split('\n'):
439
+ self.assertTrue(len(line) < width)
440
+
441
+ def test_narrow_width(self):
442
+ a = bitarray('101')
443
+ f = StringIO()
444
+ pprint(a, stream=f, width=6)
445
+ self.assertEqual(eval(f.getvalue()), a)
446
+
447
+ def test_fallback(self):
448
+ for a in None, 'asd', [1, 2], bitarray(), frozenbitarray('1'):
449
+ self.round_trip(a)
450
+
451
+ def test_subclass(self):
452
+ class Foo(bitarray):
453
+ pass
454
+
455
+ a = Foo()
456
+ code = self.get_code_string(a)
457
+ self.assertEqual(code, "Foo()\n")
458
+ b = eval(code)
459
+ self.assertEqual(b, a)
460
+ self.assertIs(type(b), type(a))
461
+
462
+ def test_random(self):
463
+ for n in range(150):
464
+ self.round_trip(urandom(n))
465
+
466
+ def test_file(self):
467
+ with tempfile.TemporaryDirectory() as tmpdir:
468
+ tmpfile = os.path.join(tmpdir, 'testfile')
469
+ a = urandom_2(1000)
470
+ with open(tmpfile, 'w') as fo:
471
+ pprint(a, fo)
472
+ with open(tmpfile, 'r') as fi:
473
+ b = eval(fi.read())
474
+ self.assertEqual(a, b)
475
+
476
+ def test_errors(self):
477
+ a = urandom(60)
478
+ self.assertRaises(TypeError, pprint, a, group=6.0)
479
+ self.assertRaises(TypeError, pprint, a, indent=0.0)
480
+ self.assertRaises(TypeError, pprint, a, width=60.0)
481
+
482
+ # ----------------------------- strip() -----------------------------------
483
+
484
+ class StripTests(unittest.TestCase, Util):
485
+
486
+ def test_simple(self):
487
+ self.assertRaises(TypeError, strip, '0110')
488
+ self.assertRaises(TypeError, strip, bitarray(), 123)
489
+ self.assertRaises(ValueError, strip, bitarray(), 'up')
490
+
491
+ a = bitarray('00010110000')
492
+ self.assertEQUAL(strip(a), bitarray('0001011'))
493
+ self.assertEQUAL(strip(a, 'left'), bitarray('10110000'))
494
+ self.assertEQUAL(strip(a, 'both'), bitarray('1011'))
495
+ b = frozenbitarray('00010110000')
496
+ c = strip(b, 'both')
497
+ self.assertEqual(c, bitarray('1011'))
498
+ self.assertIs(type(c), frozenbitarray)
499
+
500
+ def test_zeros_ones(self):
501
+ for _ in range(50):
502
+ n = randrange(10)
503
+ mode = choice(['left', 'right', 'both'])
504
+ a = zeros(n)
505
+ c = strip(a, mode)
506
+ self.assertIs(type(c), bitarray)
507
+ self.assertEqual(len(c), 0)
508
+ self.assertEqual(a, zeros(n))
509
+
510
+ b = frozenbitarray(a)
511
+ c = strip(b, mode)
512
+ self.assertIs(type(c), frozenbitarray)
513
+ self.assertEqual(len(c), 0)
514
+
515
+ a.setall(1)
516
+ c = strip(a, mode)
517
+ self.assertEqual(c, ones(n))
518
+
519
+ def test_random(self):
520
+ for a in self.randombitarrays():
521
+ b = a.copy()
522
+ f = frozenbitarray(a)
523
+ s = a.to01()
524
+ for mode, res in [
525
+ ('left', bitarray(s.lstrip('0'), a.endian)),
526
+ ('right', bitarray(s.rstrip('0'), a.endian)),
527
+ ('both', bitarray(s.strip('0'), a.endian)),
528
+ ]:
529
+ c = strip(a, mode)
530
+ self.assertEQUAL(c, res)
531
+ self.assertIs(type(c), bitarray)
532
+ self.assertEQUAL(a, b)
533
+
534
+ c = strip(f, mode)
535
+ self.assertEQUAL(c, res)
536
+ self.assertIs(type(c), frozenbitarray)
537
+ self.assertEQUAL(f, b)
538
+
539
+ def test_one_set(self):
540
+ for _ in range(10):
541
+ n = randint(1, 10000)
542
+ a = bitarray(n)
543
+ a.setall(0)
544
+ a[randrange(n)] = 1
545
+ self.assertEqual(strip(a, 'both'), bitarray('1'))
546
+ self.assertEqual(len(a), n)
547
+
548
+ # ----------------------------- count_n() ---------------------------------
549
+
550
+ class CountN_Tests(unittest.TestCase, Util):
551
+
552
+ @staticmethod
553
+ def count_n(a, n):
554
+ "return lowest index i for which a[:i].count() == n"
555
+ i, j = n, a.count(1, 0, n)
556
+ while j < n:
557
+ j += a[i]
558
+ i += 1
559
+ return i
560
+
561
+ def check_result(self, a, n, i, v=1):
562
+ self.assertEqual(a.count(v, 0, i), n)
563
+ if i == 0:
564
+ self.assertEqual(n, 0)
565
+ else:
566
+ self.assertEqual(a[i - 1], v)
567
+
568
+ def test_empty(self):
569
+ a = bitarray()
570
+ self.assertEqual(count_n(a, 0), 0)
571
+ self.assertEqual(count_n(a, 0, 0), 0)
572
+ self.assertEqual(count_n(a, 0, 1), 0)
573
+ self.assertRaises(ValueError, count_n, a, 1)
574
+ self.assertRaises(TypeError, count_n, '', 0)
575
+ self.assertRaises(TypeError, count_n, a, 7.0)
576
+ self.assertRaises(ValueError, count_n, a, 0, 2)
577
+ self.assertRaisesMessage(ValueError, "n = 1 larger than bitarray "
578
+ "length 0", count_n, a, 1)
579
+
580
+ def test_simple(self):
581
+ a = bitarray('111110111110111110111110011110111110111110111000')
582
+ b = a.copy()
583
+ self.assertEqual(len(a), 48)
584
+ self.assertEqual(a.count(), 37)
585
+ self.assertEqual(a.count(0), 11)
586
+
587
+ self.assertEqual(count_n(a, 0), 0)
588
+ self.assertEqual(count_n(a, 0, 0), 0)
589
+ self.assertEqual(count_n(a, 2, 0), 12)
590
+ self.assertEqual(count_n(a, 10, 0), 47)
591
+ self.assertEqual(count_n(a, 20), 23)
592
+ self.assertEqual(count_n(a, 20, 1), 23)
593
+ self.assertEqual(count_n(a, 37), 45)
594
+ # n < 0
595
+ self.assertRaisesMessage(ValueError, "non-negative integer expected",
596
+ count_n, a, -1)
597
+ # n > len(a)
598
+ self.assertRaisesMessage(ValueError, "n = 49 larger than bitarray "
599
+ "length 48", count_n, a, 49)
600
+ # n > a.count(0)
601
+ self.assertRaisesMessage(ValueError, "n = 12 exceeds total count "
602
+ "(a.count(0) = 11)", count_n, a, 12, 0)
603
+ # n > a.count(1)
604
+ self.assertRaisesMessage(ValueError, "n = 38 exceeds total count "
605
+ "(a.count(1) = 37)", count_n, a, 38, 1)
606
+
607
+ for v in 0, 1:
608
+ for n in range(a.count(v) + 1):
609
+ i = count_n(a, n, v)
610
+ self.check_result(a, n, i, v)
611
+ self.assertEqual(a[:i].count(v), n)
612
+ self.assertEqual(i, self.count_n(a if v else ~a, n))
613
+ self.assertEQUAL(a, b)
614
+
615
+ def test_frozenbitarray(self):
616
+ a = frozenbitarray('001111101111101111101111100111100')
617
+ self.assertEqual(len(a), 33)
618
+ self.assertEqual(a.count(), 24)
619
+ self.assertEqual(count_n(a, 0), 0)
620
+ self.assertEqual(count_n(a, 10), 13)
621
+ self.assertEqual(count_n(a, 24), 31)
622
+ self.assertRaises(ValueError, count_n, a, -1) # n < 0
623
+ self.assertRaises(ValueError, count_n, a, 25) # n > a.count()
624
+ self.assertRaises(ValueError, count_n, a, 34) # n > len(a)
625
+ for n in range(25):
626
+ self.check_result(a, n, count_n(a, n))
627
+
628
+ def test_ones(self):
629
+ n = randint(1, 100_000)
630
+ a = ones(n)
631
+ self.assertEqual(count_n(a, n), n)
632
+ self.assertRaises(ValueError, count_n, a, 1, 0)
633
+ self.assertRaises(ValueError, count_n, a, n + 1)
634
+ for _ in range(20):
635
+ i = randint(0, n)
636
+ self.assertEqual(count_n(a, i), i)
637
+
638
+ def test_one_set(self):
639
+ n = randint(1, 100_000)
640
+ a = zeros(n)
641
+ self.assertEqual(count_n(a, 0), 0)
642
+ self.assertRaises(ValueError, count_n, a, 1)
643
+ for _ in range(20):
644
+ a.setall(0)
645
+ i = randrange(n)
646
+ a[i] = 1
647
+ self.assertEqual(count_n(a, 1), i + 1)
648
+ self.assertRaises(ValueError, count_n, a, 2)
649
+
650
+ def test_last(self):
651
+ for N in range(1, 1000):
652
+ a = zeros(N)
653
+ a[-1] = 1
654
+ self.assertEqual(count_n(a, 1), N)
655
+ if N == 1:
656
+ msg = "n = 2 larger than bitarray length 1"
657
+ else:
658
+ msg = "n = 2 exceeds total count (a.count(1) = 1)"
659
+ self.assertRaisesMessage(ValueError, msg, count_n, a, 2)
660
+
661
+ def test_primes(self):
662
+ a = gen_primes(10_000)
663
+ # there are 1229 primes below 10,000
664
+ self.assertEqual(a.count(), 1229)
665
+ for n, p in [( 10, 29), # the 10th prime number is 29
666
+ ( 100, 541), # the 100th prime number is 541
667
+ (1000, 7919)]: # the 1000th prime number is 7919
668
+ self.assertEqual(count_n(a, n) - 1, p)
669
+
670
+ def test_large(self):
671
+ for _ in range(100):
672
+ N = randint(100_000, 250_000)
673
+ a = bitarray(N)
674
+ v = getrandbits(1)
675
+ a.setall(not v)
676
+ for _ in range(randrange(100)):
677
+ a[randrange(N)] = v
678
+ tc = a.count(v) # total count
679
+ i = count_n(a, tc, v)
680
+ self.check_result(a, tc, i, v)
681
+ n = tc + 1
682
+ self.assertRaisesMessage(ValueError, "n = %d exceeds total count "
683
+ "(a.count(%d) = %d)" % (n, v, tc),
684
+ count_n, a, n, v)
685
+ for _ in range(20):
686
+ n = randint(0, tc)
687
+ i = count_n(a, n, v)
688
+ self.check_result(a, n, i, v)
689
+
690
+ # ---------------------------------------------------------------------------
691
+
692
+ class BitwiseCountTests(unittest.TestCase, Util):
693
+
694
+ def test_count_byte(self):
695
+ for i in range(256):
696
+ a = bitarray(bytearray([i]))
697
+ cnt = a.count()
698
+ self.assertEqual(count_and(a, zeros(8)), 0)
699
+ self.assertEqual(count_and(a, ones(8)), cnt)
700
+ self.assertEqual(count_and(a, a), cnt)
701
+ self.assertEqual(count_or(a, zeros(8)), cnt)
702
+ self.assertEqual(count_or(a, ones(8)), 8)
703
+ self.assertEqual(count_or(a, a), cnt)
704
+ self.assertEqual(count_xor(a, zeros(8)), cnt)
705
+ self.assertEqual(count_xor(a, ones(8)), 8 - cnt)
706
+ self.assertEqual(count_xor(a, a), 0)
707
+
708
+ def test_1(self):
709
+ a = bitarray('001111')
710
+ aa = a.copy()
711
+ b = bitarray('010011')
712
+ bb = b.copy()
713
+ self.assertEqual(count_and(a, b), 2)
714
+ self.assertEqual(count_or(a, b), 5)
715
+ self.assertEqual(count_xor(a, b), 3)
716
+ for f in count_and, count_or, count_xor:
717
+ # not two arguments
718
+ self.assertRaises(TypeError, f)
719
+ self.assertRaises(TypeError, f, a)
720
+ self.assertRaises(TypeError, f, a, b, 3)
721
+ # wrong argument types
722
+ self.assertRaises(TypeError, f, a, '')
723
+ self.assertRaises(TypeError, f, '1', b)
724
+ self.assertRaises(TypeError, f, a, 4)
725
+ self.assertEQUAL(a, aa)
726
+ self.assertEQUAL(b, bb)
727
+
728
+ b.append(1)
729
+ for f in count_and, count_or, count_xor:
730
+ self.assertRaises(ValueError, f, a, b)
731
+ self.assertRaises(ValueError, f,
732
+ bitarray('110', 'big'),
733
+ bitarray('101', 'little'))
734
+
735
+ def test_frozen(self):
736
+ a = frozenbitarray('001111')
737
+ b = frozenbitarray('010011')
738
+ self.assertEqual(count_and(a, b), 2)
739
+ self.assertEqual(count_or(a, b), 5)
740
+ self.assertEqual(count_xor(a, b), 3)
741
+
742
+ def test_random(self):
743
+ for _ in range(100):
744
+ n = randrange(1000)
745
+ a = urandom_2(n)
746
+ b = urandom(n, a.endian)
747
+ self.assertEqual(count_and(a, b), (a & b).count())
748
+ self.assertEqual(count_or(a, b), (a | b).count())
749
+ self.assertEqual(count_xor(a, b), (a ^ b).count())
750
+
751
+ def test_misc(self):
752
+ for a in self.randombitarrays():
753
+ n = len(a)
754
+ b = urandom(n, a.endian)
755
+ # any and
756
+ self.assertEqual(any(a & b), count_and(a, b) > 0)
757
+ self.assertEqual(any_and(a, b), any(a & b))
758
+ # any or
759
+ self.assertEqual(any(a | b), count_or(a, b) > 0)
760
+ self.assertEqual(any(a | b), any(a) or any(b))
761
+ # any xor
762
+ self.assertEqual(any(a ^ b), count_xor(a, b) > 0)
763
+ self.assertEqual(any(a ^ b), a != b)
764
+
765
+ # all and
766
+ self.assertEqual(all(a & b), count_and(a, b) == n)
767
+ self.assertEqual(all(a & b), all(a) and all(b))
768
+ # all or
769
+ self.assertEqual(all(a | b), count_or(a, b) == n)
770
+ # all xor
771
+ self.assertEqual(all(a ^ b), count_xor(a, b) == n)
772
+ self.assertEqual(all(a ^ b), a == ~b)
773
+
774
+ # --------------------------- any_and() -----------------------------------
775
+
776
+ class BitwiseAnyTests(unittest.TestCase, Util):
777
+
778
+ def test_basic(self):
779
+ a = frozenbitarray('0101')
780
+ b = bitarray('0111')
781
+ self.assertTrue(any_and(a, b))
782
+ self.assertRaises(TypeError, any_and)
783
+ self.assertRaises(TypeError, any_and, a, 4)
784
+ b.append(1)
785
+ self.assertRaises(ValueError, any_and, a, b)
786
+ self.assertRaises(ValueError, any_and,
787
+ bitarray('01', 'little'),
788
+ bitarray('11', 'big'))
789
+
790
+ def test_overlap(self):
791
+ n = 100
792
+ for _ in range(500):
793
+ i1 = randint(0, n)
794
+ j1 = randint(i1, n)
795
+ r1 = range(i1, j1)
796
+
797
+ i2 = randint(0, n)
798
+ j2 = randint(i2, n)
799
+ r2 = range(i2, j2)
800
+
801
+ # test if ranges r1 and r2 overlap
802
+ res1 = bool(r1) and bool(r2) and (i2 in r1 or i1 in r2)
803
+ res2 = bool(set(r1) & set(r2))
804
+ self.assertEqual(res1, res2)
805
+
806
+ a1, a2 = bitarray(n), bitarray(n)
807
+ a1[i1:j1] = a2[i2:j2] = 1
808
+ self.assertEqual(any_and(a1, a2), res1)
809
+
810
+ def test_common(self):
811
+ n = 100
812
+ for _ in range(500):
813
+ s1 = self.random_slice(n)
814
+ s2 = self.random_slice(n)
815
+ r1 = range(n)[s1]
816
+ r2 = range(n)[s2]
817
+ # test if ranges r1 and r2 have common items
818
+ a1, a2 = bitarray(n), bitarray(n)
819
+ a1[s1] = a2[s2] = 1
820
+ self.assertEqual(any_and(a1, a2), bool(set(r1) & set(r2)))
821
+
822
+ def check(self, a, b):
823
+ r = any_and(a, b)
824
+ self.assertIs(type(r), bool)
825
+ self.assertEqual(r, any_and(b, a)) # symmetry
826
+ self.assertEqual(r, any(a & b))
827
+ self.assertEqual(r, (a & b).any())
828
+ self.assertEqual(r, count_and(a, b) > 0)
829
+
830
+ def test_explicit(self):
831
+ for a, b , res in [
832
+ ('', '', False),
833
+ ('0', '1', False),
834
+ ('0', '0', False),
835
+ ('1', '1', True),
836
+ ('00011', '11100', False),
837
+ ('00001011 1', '01000100 1', True)]:
838
+ a = bitarray(a)
839
+ b = bitarray(b)
840
+ self.assertIs(any_and(a, b), res)
841
+ self.check(a, b)
842
+
843
+ def test_random(self):
844
+ for a in self.randombitarrays():
845
+ n = len(a)
846
+ b = urandom(n, a.endian)
847
+ self.check(a, b)
848
+
849
+ def test_one(self):
850
+ for n in range(1, 300):
851
+ a = zeros(n)
852
+ b = urandom(n)
853
+ i = randrange(n)
854
+ a[i] = 1
855
+ self.assertEqual(b[i], any_and(a, b))
856
+
857
+ # ---------------------------- subset() -----------------------------------
858
+
859
+ class SubsetTests(unittest.TestCase, Util):
860
+
861
+ def test_basic(self):
862
+ a = frozenbitarray('0101')
863
+ b = bitarray('0111')
864
+ self.assertTrue(subset(a, b))
865
+ self.assertFalse(subset(b, a))
866
+ self.assertRaises(TypeError, subset)
867
+ self.assertRaises(TypeError, subset, a, '')
868
+ self.assertRaises(TypeError, subset, '1', b)
869
+ self.assertRaises(TypeError, subset, a, 4)
870
+ b.append(1)
871
+ self.assertRaises(ValueError, subset, a, b)
872
+ self.assertRaises(ValueError, subset,
873
+ bitarray('01', 'little'),
874
+ bitarray('11', 'big'))
875
+
876
+ def check(self, a, b, res):
877
+ r = subset(a, b)
878
+ self.assertIs(type(r), bool)
879
+ self.assertEqual(r, res)
880
+ self.assertEqual(a | b == b, res)
881
+ self.assertEqual(a & b == a, res)
882
+
883
+ def test_true(self):
884
+ for a, b in [('', ''), ('0', '1'), ('0', '0'), ('1', '1'),
885
+ ('000', '111'), ('0101', '0111'),
886
+ ('000010111', '010011111')]:
887
+ self.check(bitarray(a), bitarray(b), True)
888
+
889
+ def test_false(self):
890
+ for a, b in [('1', '0'), ('1101', '0111'),
891
+ ('0000101111', '0100111011')]:
892
+ self.check(bitarray(a), bitarray(b), False)
893
+
894
+ def test_random(self):
895
+ for a in self.randombitarrays(start=1):
896
+ b = a.copy()
897
+ # we set one random bit in b to 1, so a is always a subset of b
898
+ b[randrange(len(a))] = 1
899
+ self.check(a, b, True)
900
+ # but b is only a subset when they are equal
901
+ self.check(b, a, a == b)
902
+ # we set all bits in a, which ensures that b is a subset of a
903
+ a.setall(1)
904
+ self.check(b, a, True)
905
+
906
+ # ------------------------- correspond_all() ------------------------------
907
+
908
+ class CorrespondAllTests(unittest.TestCase):
909
+
910
+ def test_basic(self):
911
+ a = frozenbitarray('0101')
912
+ b = bitarray('0110')
913
+ self.assertEqual(correspond_all(a, b), (1, 1, 1, 1))
914
+ self.assertRaises(TypeError, correspond_all)
915
+ b.append(1)
916
+ self.assertRaises(ValueError, correspond_all, a, b)
917
+ self.assertRaises(ValueError, correspond_all,
918
+ bitarray('01', 'little'),
919
+ bitarray('11', 'big'))
920
+
921
+ def test_explicit(self):
922
+ for a, b, res in [
923
+ ('', '', (0, 0, 0, 0)),
924
+ ('0000011111',
925
+ '0000100111', (4, 1, 2, 3)),
926
+ ]:
927
+ self.assertEqual(correspond_all(bitarray(a), bitarray(b)), res)
928
+
929
+ def test_random(self):
930
+ for _ in range(100):
931
+ n = randrange(3000)
932
+ a = urandom_2(n)
933
+ b = urandom(n, a.endian)
934
+ res = correspond_all(a, b)
935
+ self.assertEqual(res[0], count_and(~a, ~b))
936
+ self.assertEqual(res[1], count_and(~a, b))
937
+ self.assertEqual(res[2], count_and(a, ~b))
938
+ self.assertEqual(res[3], count_and(a, b))
939
+
940
+ self.assertEqual(res[0], n - count_or(a, b))
941
+ self.assertEqual(res[1] + res[2], count_xor(a, b))
942
+ self.assertEqual(sum(res), n)
943
+
944
+ # ----------------------------- byteswap() --------------------------------
945
+
946
+ @unittest.skipIf(is_pypy, "skip test on PyPy")
947
+ class ByteSwapTests(unittest.TestCase):
948
+
949
+ def test_basic_bytearray(self):
950
+ a = bytearray(b"ABCD")
951
+ byteswap(a, 2)
952
+ self.assertEqual(a, bytearray(b"BADC"))
953
+ byteswap(a)
954
+ self.assertEqual(a, bytearray(b"CDAB"))
955
+
956
+ a = bytearray(b"ABCDEF")
957
+ byteswap(a, 3)
958
+ self.assertEqual(a, bytearray(b"CBAFED"))
959
+ byteswap(a, 1)
960
+ self.assertEqual(a, bytearray(b"CBAFED"))
961
+
962
+ def test_basic_bitarray(self):
963
+ a = bitarray("11110000 01010101")
964
+ byteswap(a)
965
+ self.assertEqual(a, bitarray("01010101 11110000"))
966
+
967
+ a = bitarray("01111000 1001")
968
+ b = a.copy()
969
+ a.tobytes() # clear padbits
970
+ byteswap(a)
971
+ self.assertEqual(a, bitarray("10010000 0111"))
972
+ byteswap(a)
973
+ self.assertEqual(a, b)
974
+
975
+ def test_basic_array(self):
976
+ r = os.urandom(64)
977
+ for typecode in array.typecodes:
978
+ # type code 'u' is deprecated and will be removed in Python 3.16
979
+ if typecode in ('u', 'Zf', 'Zd'):
980
+ continue
981
+ a = array.array(typecode, r)
982
+ self.assertEqual(len(a) * a.itemsize, 64)
983
+ a.byteswap()
984
+ byteswap(a, a.itemsize)
985
+ self.assertEqual(a.tobytes(), r)
986
+
987
+ def test_empty(self):
988
+ a = bytearray()
989
+ byteswap(a)
990
+ self.assertEqual(a, bytearray())
991
+ for n in range(10):
992
+ byteswap(a, n)
993
+ self.assertEqual(a, bytearray())
994
+
995
+ def test_one_byte(self):
996
+ a = bytearray(b'\xab')
997
+ byteswap(a)
998
+ self.assertEqual(a, bytearray(b'\xab'))
999
+ for n in range(2):
1000
+ byteswap(a, n)
1001
+ self.assertEqual(a, bytearray(b'\xab'))
1002
+
1003
+ def test_errors(self):
1004
+ # buffer not writable
1005
+ for a in b"AB", frozenbitarray(16):
1006
+ self.assertRaises(BufferError, byteswap, a)
1007
+
1008
+ a = bytearray(b"ABCD")
1009
+ b = bitarray(32)
1010
+ for n in -1, 3, 5, 6:
1011
+ # byte size not multiple of n
1012
+ self.assertRaises(ValueError, byteswap, a, n)
1013
+ self.assertRaises(ValueError, byteswap, b, n)
1014
+
1015
+ def test_range(self):
1016
+ for n in range(20):
1017
+ for m in range(20):
1018
+ r = os.urandom(m * n)
1019
+ a = bytearray(r)
1020
+ byteswap(a, n)
1021
+ lst = []
1022
+ for i in range(m):
1023
+ x = r[i * n:i * n + n]
1024
+ lst.extend(x[::-1])
1025
+ self.assertEqual(a, bytearray(lst))
1026
+
1027
+ def test_reverse_bytearray(self):
1028
+ for n in range(100):
1029
+ r = os.urandom(n)
1030
+ a = bytearray(r)
1031
+ byteswap(a)
1032
+ self.assertEqual(a, bytearray(r[::-1]))
1033
+
1034
+ def test_reverse_bitarray(self):
1035
+ for n in range(100):
1036
+ a = urandom(8 * n)
1037
+ b = a.copy()
1038
+ byteswap(a)
1039
+ a.bytereverse()
1040
+ self.assertEqual(a, b[::-1])
1041
+
1042
+ # ------------------------------ parity() ---------------------------------
1043
+
1044
+ class ParityTests(unittest.TestCase):
1045
+
1046
+ def test_explicit(self):
1047
+ for s, res in [('', 0), ('1', 1), ('0010011', 1), ('10100110', 0)]:
1048
+ self.assertIs(parity(bitarray(s)), res)
1049
+ self.assertIs(parity(frozenbitarray(s)), res)
1050
+
1051
+ def test_zeros_ones(self):
1052
+ for n in range(2000):
1053
+ self.assertEqual(parity(zeros(n)), 0)
1054
+ self.assertEqual(parity(ones(n)), n % 2)
1055
+
1056
+ def test_random(self):
1057
+ endian = choice(ENDIANS)
1058
+ a = bitarray(endian=endian)
1059
+ par = 0
1060
+ for i in range(2000):
1061
+ self.assertEqual(parity(a), par)
1062
+ self.assertEqual(par, a.count() % 2)
1063
+ self.assertEqual(a.endian, endian)
1064
+ self.assertEqual(len(a), i)
1065
+ v = getrandbits(1)
1066
+ a.append(v)
1067
+ par ^= v
1068
+
1069
+ def test_wrong_args(self):
1070
+ self.assertRaises(TypeError, parity, '')
1071
+ self.assertRaises(TypeError, parity, 1)
1072
+ self.assertRaises(TypeError, parity)
1073
+ self.assertRaises(TypeError, parity, bitarray("110"), 1)
1074
+
1075
+ # ---------------------------- sum_indices() ------------------------------
1076
+
1077
+ class SumIndicesTests(unittest.TestCase):
1078
+
1079
+ # Additional tests in: devel/test_sum_indices.py
1080
+
1081
+ def test_explicit(self):
1082
+ endian = choice(ENDIANS)
1083
+ a = bitarray("01100111 1101", endian)
1084
+ self.assertEqual(sum_indices(a, mode=1), 49)
1085
+ self.assertEqual(sum_indices(a, mode=2), 381)
1086
+ a = frozenbitarray("00001111", endian)
1087
+ self.assertEqual(sum_indices(a, mode=1), 22)
1088
+ self.assertEqual(sum_indices(a, mode=2), 126)
1089
+
1090
+ def test_wrong_args(self):
1091
+ S = sum_indices
1092
+ self.assertRaises(TypeError, S, '')
1093
+ self.assertRaises(TypeError, S, 1.0)
1094
+ self.assertRaises(TypeError, S)
1095
+ for mode in -1, 0, 3, 4:
1096
+ self.assertRaises(ValueError, S, bitarray("110"), mode)
1097
+
1098
+ def test_ones(self):
1099
+ a = bitarray(endian=choice(ENDIANS))
1100
+ sm1 = sm2 = 0
1101
+ for i in range(100):
1102
+ a.append(1)
1103
+ sm1 += i
1104
+ sm2 += i * i
1105
+ self.assertEqual(sum_indices(a, 1), sm1)
1106
+ self.assertEqual(sum_indices(a, 2), sm2)
1107
+
1108
+ def check_sparse(self, n, k, mode):
1109
+ a = zeros(n, choice(ENDIANS))
1110
+ indices = sample(range(n), k)
1111
+ a[indices] = 1
1112
+ res = sum(indices) if mode == 1 else sum(i * i for i in indices)
1113
+ self.assertEqual(a.count(), k)
1114
+ self.assertEqual(sum_indices(a, mode), res)
1115
+
1116
+ def test_small(self):
1117
+ for mode in 1, 2:
1118
+ for a in bitarray(), bitarray("0"), bitarray("1"):
1119
+ self.assertEqual(sum_indices(a, mode), 0)
1120
+
1121
+ for n in range(2, 50):
1122
+ k = randrange(n // 2)
1123
+ self.check_sparse(n, k, mode)
1124
+
1125
+ def test_sparse(self):
1126
+ for _ in range(10):
1127
+ n = randrange(2, 1_500_000) # below and above block size
1128
+ k = randrange(min(100, n // 2))
1129
+ mode = randint(1, 2)
1130
+ self.check_sparse(n, k, mode)
1131
+
1132
+ # ---------------------------------------------------------------------------
1133
+
1134
+ class XoredIndicesTests(unittest.TestCase, Util):
1135
+
1136
+ def test_explicit(self):
1137
+ for s, r in [("", 0), ("0", 0), ("1", 0), ("11", 1),
1138
+ ("011", 3), ("001", 2), ("0001100", 7),
1139
+ ("01100111 1101", 13)]:
1140
+ for a in [bitarray(s, choice(ENDIANS)),
1141
+ frozenbitarray(s, choice(ENDIANS))]:
1142
+ self.assertEqual(xor_indices(a), r)
1143
+
1144
+ def test_wrong_args(self):
1145
+ X = xor_indices
1146
+ self.assertRaises(TypeError, X, '')
1147
+ self.assertRaises(TypeError, X, 1)
1148
+ self.assertRaises(TypeError, X)
1149
+ self.assertRaises(TypeError, X, bitarray("110"), 1)
1150
+
1151
+ def test_ones(self):
1152
+ # OEIS A003815
1153
+ lst = [0, 1, 3, 0, 4, 1, 7, 0, 8, 1, 11, 0, 12, 1, 15, 0, 16, 1, 19]
1154
+ self.assertEqual([xor_indices(ones(i)) for i in range(1, 20)], lst)
1155
+ a = bitarray()
1156
+ x = 0
1157
+ for i in range(1000):
1158
+ a.append(1)
1159
+ x ^= i
1160
+ self.assertEqual(xor_indices(a), x)
1161
+ if i < 19:
1162
+ self.assertEqual(lst[i], x)
1163
+
1164
+ def test_primes(self):
1165
+ # OEIS A126084
1166
+ lst = [0, 2, 1, 4, 3, 8, 5, 20, 7, 16, 13, 18, 55, 30, 53, 26, 47]
1167
+ primes = gen_primes(1000)
1168
+ x = 0
1169
+ for i, p in enumerate(primes.search(1)):
1170
+ self.assertEqual(xor_indices(primes[:p]), x)
1171
+ if i < 17:
1172
+ self.assertEqual(lst[i], x)
1173
+ x ^= p
1174
+
1175
+ def test_large_random(self):
1176
+ n = 10_037
1177
+ for a in [urandom_2(n), frozenbitarray(urandom_2(n))]:
1178
+ res = reduce(operator.xor, (i for i, v in enumerate(a) if v))
1179
+ b = a.copy()
1180
+ self.assertEqual(xor_indices(a), res)
1181
+ self.assertEqual(a, b)
1182
+
1183
+ def test_random(self):
1184
+ for a in self.randombitarrays():
1185
+ c = 0
1186
+ for i, v in enumerate(a):
1187
+ c ^= i * v
1188
+ self.assertEqual(xor_indices(a), c)
1189
+
1190
+ def test_flips(self):
1191
+ a = bitarray(128)
1192
+ c = 0
1193
+ for _ in range(1000):
1194
+ self.assertEqual(xor_indices(a), c)
1195
+ i = randrange(len(a))
1196
+ a.invert(i)
1197
+ c ^= i
1198
+
1199
+ def test_error_correct(self):
1200
+ parity_bits = [1, 2, 4, 8, 16, 32, 64, 128] # parity bit positions
1201
+ a = urandom(256)
1202
+ a[parity_bits] = 0
1203
+ c = xor_indices(a)
1204
+ # set parity bits such that block is well prepared
1205
+ a[parity_bits] = int2ba(c, length=8, endian="little")
1206
+ for i in range(0, 256):
1207
+ self.assertEqual(xor_indices(a), 0) # ensure well prepared
1208
+ a.invert(i)
1209
+ self.assertEqual(xor_indices(a), i) # index of the flipped bit!
1210
+ a.invert(i)
1211
+
1212
+
1213
+ # ------------------------- run-length codec ----------------------------
1214
+
1215
+ class RL_Util:
1216
+
1217
+ @staticmethod
1218
+ def runs(a):
1219
+ "return number of uninterrupted intervals of 1s and 0s"
1220
+ n = len(a)
1221
+ if n < 2:
1222
+ return n
1223
+ return count_xor(a[:-1], a[1:]) + 1
1224
+
1225
+ def random_runs(self, n, k):
1226
+ "generate random bitarray of length `n` with `k` runs."
1227
+ self.assertTrue(n == k == 0 or 1 <= k <= n)
1228
+ a = bitarray(n)
1229
+ if n == 0:
1230
+ return a
1231
+ vi = getrandbits(1)
1232
+
1233
+ # Create sorted sample from `range(1, n)`, with `n` last and
1234
+ # with a total of `k` items.
1235
+ lst = sample(range(1, n), k - 1)
1236
+ lst.append(n)
1237
+ lst.sort()
1238
+
1239
+ start = 0
1240
+ for stop in lst:
1241
+ a[start:stop] = vi
1242
+ vi = not vi
1243
+ start = stop
1244
+ return a
1245
+
1246
+
1247
+ class RL_Tests(unittest.TestCase, Util, RL_Util):
1248
+
1249
+ def test_explicit(self):
1250
+ for s, b in [
1251
+ ("", b'0\x00'),
1252
+ ("0", b'0\x01\x01'),
1253
+ ("1", b'1\x01\x01'),
1254
+ ("10", b'1\x02\x01\x01'),
1255
+ ("01", b'0\x02\x01\x01'),
1256
+ ("1" + 63 * "0", b'1\x40\x01\x3f'),
1257
+ ("0" + 63 * "1", b'0\x40\x01\x3f'),
1258
+ ]:
1259
+ a = bitarray(s)
1260
+ self.assertEqual(rl_encode(a), b)
1261
+ self.assertEqual(rl_decode(b), a)
1262
+ it = iter(b)
1263
+ self.assertEqual(rl_decode(it), a)
1264
+ self.assertRaises(StopIteration, next, it)
1265
+ it = iter(b + b'XYZ')
1266
+ self.assertEqual(rl_decode(it), a)
1267
+ self.assertEqual(next(it), ord(b'X'))
1268
+
1269
+ def test_decode_errors(self):
1270
+ # incomplete stream
1271
+ for b in [
1272
+ b'', # no first bit
1273
+ b'0', # no uleb128 encoded nbits
1274
+ b'1\x80', # nbits uleb128 encoding incomplete
1275
+ b'0\x01', # run length incomplete
1276
+ b'1\x02\x01', # run length incomplete
1277
+ ]:
1278
+ self.assertRaises(StopIteration, rl_decode, b)
1279
+
1280
+ # invalid first bit
1281
+ for b in b'\x00', b'\x01', b'2':
1282
+ msg = "invalid first bit: 0x%02x" % ord(b)
1283
+ self.assertRaisesMessage(ValueError, msg, rl_decode, b)
1284
+
1285
+ # run length cannot be zero
1286
+ for b in b'0\x01\x00\x01', b'1\xff\x0a\x01\x00':
1287
+ msg = "run length cannot be zero"
1288
+ self.assertRaisesMessage(ValueError, msg, rl_decode, b)
1289
+
1290
+ # sequence exceeds nbits
1291
+ for b, nrun, nrem in [
1292
+ (b'1\x01\x02', 2, 1),
1293
+ (b'1\x30\x2f\x02', 2, 1),
1294
+ (b'0\x40\x01\x40', 64, 63),
1295
+ ]:
1296
+ msg = "run length %d exceeds remaining length %d" % (nrun, nrem)
1297
+ self.assertRaisesMessage(ValueError, msg, rl_decode, b)
1298
+
1299
+ def test_decode_ambiguity(self):
1300
+ # For an empty bitarray either first-bit byte works.
1301
+ for b in b'0\x00', b'1\x00':
1302
+ self.assertEqual(rl_decode(b), bitarray())
1303
+
1304
+ # ULEB128 permits overlong representations.
1305
+ for b in b'0\x00', b'0\x80\x00':
1306
+ self.assertEqual(rl_decode(b), bitarray())
1307
+
1308
+ def test_decode_endian(self):
1309
+ b = b'0\x02\x01\x01'
1310
+ for endian in OPT_ENDIANS:
1311
+ a = rl_decode(b, endian=endian)
1312
+ self.assertEqual(a.endian, endian or get_default_endian())
1313
+ self.assertEqual(a, bitarray("01"))
1314
+
1315
+ def test_decode_types(self):
1316
+ b = b'0\x02\x01\x01'
1317
+ for c in b, bytearray(b), list(b), iter(b):
1318
+ a = rl_decode(c)
1319
+ self.assertIs(type(a), bitarray)
1320
+ self.assertEqual(a.to01(), '01')
1321
+
1322
+ def test_encode_types(self):
1323
+ a = frozenbitarray("00001111 11111111 11110000 000")
1324
+ self.assertEqual(len(a), 0x1b)
1325
+ b = rl_encode(a)
1326
+ self.assertIs(type(b), bytes)
1327
+ self.assertEqual(b, b'0\x1b\x04\x10\x07')
1328
+
1329
+ def test_intervals(self):
1330
+ a = bitarray("00000011111111111111110000000111000000")
1331
+ b = bytearray([ord("0")])
1332
+ b.extend(uleb128_encode(len(a)))
1333
+ for unused_value, start, stop in intervals(a):
1334
+ b.extend(uleb128_encode(stop - start))
1335
+ self.assertEqual(rl_decode(b), a)
1336
+ self.assertEqual(b, rl_encode(a))
1337
+
1338
+ def test_concatenate(self):
1339
+ N = 100
1340
+ arrays = [urandom(randrange(5)) for _ in range(N)]
1341
+ it = iter(b''.join(rl_encode(a) for a in arrays))
1342
+ self.assertEqual([rl_decode(it) for _ in range(N)], arrays)
1343
+ self.assertRaises(StopIteration, next, it)
1344
+
1345
+ def test_random(self):
1346
+ for _ in range(10):
1347
+ n = randrange(1, 100)
1348
+ k = randint(1, max(1, n // 2))
1349
+ a = self.random_runs(n, k)
1350
+
1351
+ b = bytearray([ord("0") + a[0]])
1352
+ b.extend(uleb128_encode(len(a)))
1353
+ for _, start, stop in intervals(a):
1354
+ b.extend(uleb128_encode(stop - start))
1355
+
1356
+ self.assertEqual(rl_decode(b), a)
1357
+ self.assertEqual(b, rl_encode(a))
1358
+
1359
+
1360
+ # ------------------ intervals of uninterrupted runs --------------------
1361
+
1362
+ class IntervalsTests(unittest.TestCase, Util, RL_Util):
1363
+
1364
+ def test_explicit(self):
1365
+ for s, lst in [
1366
+ ('', []),
1367
+ ('0', [(0, 0, 1)]),
1368
+ ('1', [(1, 0, 1)]),
1369
+ ('00111100 0000011',
1370
+ [(0, 0, 2), (1, 2, 6), (0, 6, 13), (1, 13, 15)]),
1371
+ ]:
1372
+ a = bitarray(s)
1373
+ self.assertEqual(list(intervals(a)), lst)
1374
+ self.assertEqual(self.runs(a), len(lst))
1375
+
1376
+ def test_uniform(self):
1377
+ for n in range(1, 100):
1378
+ a = bitarray(n, choice(ENDIANS))
1379
+ for v in 0, 1:
1380
+ a.setall(v)
1381
+ self.assertEqual(list(intervals(a)), [(v, 0, n)])
1382
+ self.assertEqual(self.runs(a), 1)
1383
+
1384
+ def test_random(self):
1385
+ for _ in range(10):
1386
+ n = randrange(2, 500)
1387
+ k = randint(1, n // 2)
1388
+ a = self.random_runs(n, k)
1389
+
1390
+ # Reconstruct `a` from its intervals, starting with random data.
1391
+ b = urandom(n)
1392
+ for value, start, stop in intervals(a):
1393
+ self.assertIs(type(value), int)
1394
+ self.assertTrue(0 <= start < stop <= n)
1395
+ b[start:stop] = value
1396
+ self.assertEqual(a, b)
1397
+
1398
+ def test_list_runs(self):
1399
+ for a in self.randombitarrays():
1400
+ # list of length of runs of alternating bits
1401
+ alt_runs = [stop - start for _, start, stop in intervals(a)]
1402
+ self.assertEqual(len(alt_runs), self.runs(a))
1403
+
1404
+ b = bitarray()
1405
+ v = a[0] if a else None # value of first run
1406
+ for length in alt_runs:
1407
+ self.assertTrue(length > 0)
1408
+ b.extend(length * bitarray([v]))
1409
+ v = not v
1410
+ self.assertEqual(a, b)
1411
+
1412
+
1413
+ # -------------------------- ba2hex() hex2ba() ---------------------------
1414
+
1415
+ class HexlifyTests(unittest.TestCase, Util):
1416
+
1417
+ def test_explicit(self):
1418
+ data = [ # little big
1419
+ ('', '', ''),
1420
+ ('1000', '1', '8'),
1421
+ ('0101 0110', 'a6', '56'),
1422
+ ('0100 1001 1101', '29b', '49d'),
1423
+ ('0000 1100 1110 1111', '037f', '0cef'),
1424
+ ]
1425
+ for bs, hex_le, hex_be in data:
1426
+ a_be = bitarray(bs, 'big')
1427
+ a_le = bitarray(bs, 'little')
1428
+ self.assertEQUAL(hex2ba(hex_be, 'big'), a_be)
1429
+ self.assertEQUAL(hex2ba(hex_le, 'little'), a_le)
1430
+ self.assertEqual(ba2hex(a_be), hex_be)
1431
+ self.assertEqual(ba2hex(a_le), hex_le)
1432
+
1433
+ def test_ba2hex_group(self):
1434
+ a = bitarray('1000 0000 0101 1111', 'little')
1435
+ self.assertEqual(ba2hex(a), "10af")
1436
+ self.assertEqual(ba2hex(a, 0), "10af")
1437
+ self.assertEqual(ba2hex(a, 1, ""), "10af")
1438
+ self.assertEqual(ba2hex(a, 1), "1 0 a f")
1439
+ self.assertEqual(ba2hex(a, group=2), "10 af")
1440
+ self.assertEqual(ba2hex(a, 2, "-"), "10-af")
1441
+ self.assertEqual(ba2hex(a, group=3, sep="_"), "10a_f")
1442
+ self.assertEqual(ba2hex(a, 3, sep=", "), "10a, f")
1443
+
1444
+ def test_ba2hex_errors(self):
1445
+ self.assertRaises(TypeError, ba2hex)
1446
+ self.assertRaises(TypeError, ba2hex, None)
1447
+ self.assertRaises(TypeError, ba2hex, '101')
1448
+
1449
+ # length not multiple of 4
1450
+ self.assertRaises(ValueError, ba2hex, bitarray('10'))
1451
+
1452
+ a = bitarray('1000 0000 0101 1111', 'little')
1453
+ self.assertRaises(ValueError, ba2hex, a, -1)
1454
+ self.assertRaises(ValueError, ba2hex, a, group=-1)
1455
+ # sep not str
1456
+ self.assertRaises(TypeError, ba2hex, a, 1, b" ")
1457
+ # embedded null character in sep
1458
+ self.assertRaises(ValueError, ba2hex, a, 2, " \0")
1459
+
1460
+ def test_hex2ba_whitespace(self):
1461
+ self.assertEqual(hex2ba("F1 FA %s f3 c0" % whitespace),
1462
+ bitarray("11110001 11111010 11110011 11000000"
1463
+ if get_default_endian() == "big" else
1464
+ "11111000 11110101 11111100 00110000"))
1465
+ self.assertEQUAL(hex2ba(b' a F ', 'big'),
1466
+ bitarray('1010 1111', 'big'))
1467
+ self.assertEQUAL(hex2ba(860 * " " + '0 1D' + 590 * " ", 'little'),
1468
+ bitarray('0000 1000 1011', 'little'))
1469
+ self.assertRaises(ValueError, hex2ba, "f3\u00a0") # no-break space
1470
+ self.assertRaises(ValueError, hex2ba, b"f3\xa0")
1471
+
1472
+ def test_hex2ba_errors(self):
1473
+ self.assertRaises(TypeError, hex2ba, 0)
1474
+ self.assertRaises(TypeError, hex2ba, "F", 1)
1475
+ self.assertRaises(ValueError, hex2ba, "F", "foo")
1476
+
1477
+ for s in '01a7g89', '0\u20ac', '0 \0', b'\x00':
1478
+ self.assertRaises(ValueError, hex2ba, s)
1479
+
1480
+ for s in 'g', 'ag', 'aag', 'aaaga', 'ag':
1481
+ msg = "invalid digit found for base16, got 'g' (0x67)"
1482
+ self.assertRaisesMessage(ValueError, msg, hex2ba, s, 'big')
1483
+
1484
+ def test_hex2ba_types(self):
1485
+ for c in 'e', 'E', b'e', b'E', bytearray(b'e'), bytearray(b'E'):
1486
+ a = hex2ba(c, "big")
1487
+ self.assertEqual(a.to01(), '1110')
1488
+ self.assertEqual(a.endian, 'big')
1489
+ self.assertIs(type(a), bitarray)
1490
+
1491
+ def test_random(self):
1492
+ for _ in range(100):
1493
+ endian = choice(OPT_ENDIANS)
1494
+ a = urandom(4 * randrange(100), endian)
1495
+ s = ba2hex(a, group=randrange(10), sep=choice(whitespace))
1496
+ b = hex2ba(s, endian)
1497
+ self.assertEqual(b.endian, endian or get_default_endian())
1498
+ self.assertEqual(a, b)
1499
+ self.check_obj(b)
1500
+
1501
+ def test_hexdigits(self):
1502
+ a = hex2ba(hexdigits)
1503
+ self.assertEqual(len(a), 4 * len(hexdigits))
1504
+ self.assertIs(type(a), bitarray)
1505
+ self.check_obj(a)
1506
+
1507
+ t = ba2hex(a)
1508
+ self.assertEqual(t, hexdigits.lower())
1509
+ self.assertIs(type(t), str)
1510
+ self.assertEQUAL(a, hex2ba(t))
1511
+
1512
+ def test_binascii(self):
1513
+ a = urandom(80, 'big')
1514
+ s = binascii.hexlify(a.tobytes()).decode()
1515
+ self.assertEqual(ba2hex(a), s)
1516
+ b = bitarray(binascii.unhexlify(s), endian='big')
1517
+ self.assertEQUAL(hex2ba(s, 'big'), b)
1518
+
1519
+ # -------------------------- ba2base() base2ba() -------------------------
1520
+
1521
+ class BaseTests(unittest.TestCase, Util):
1522
+
1523
+ def test_explicit(self):
1524
+ data = [ # n little big
1525
+ ('', 2, '', ''),
1526
+ ('1 0 1', 2, '101', '101'),
1527
+ ('11 01 00', 4, '320', '310'),
1528
+ ('111 001', 8, '74', '71'),
1529
+ ('1111 0001', 16, 'f8', 'f1'),
1530
+ ('11111 00001', 32, '7Q', '7B'),
1531
+ ('111111 000001', 64, '/g', '/B'),
1532
+ ]
1533
+ for bs, n, s_le, s_be in data:
1534
+ a_le = bitarray(bs, 'little')
1535
+ a_be = bitarray(bs, 'big')
1536
+ self.assertEQUAL(base2ba(n, s_le, 'little'), a_le)
1537
+ self.assertEQUAL(base2ba(n, s_be, 'big'), a_be)
1538
+ self.assertEqual(ba2base(n, a_le), s_le)
1539
+ self.assertEqual(ba2base(n, a_be), s_be)
1540
+
1541
+ def test_base2ba_types(self):
1542
+ for c in '7', b'7', bytearray(b'7'):
1543
+ a = base2ba(32, c)
1544
+ self.assertEqual(a.to01(), '11111')
1545
+ self.assertIs(type(a), bitarray)
1546
+
1547
+ def test_base2ba_whitespace(self):
1548
+ self.assertEqual(base2ba(8, bytearray(b"17 0"), "little"),
1549
+ bitarray("100 111 000"))
1550
+ self.assertEqual(base2ba(32, "7 A"), bitarray("11111 00000"))
1551
+ self.assertEqual(base2ba(64, b"A /"), bitarray("000000 111111"))
1552
+ for n in 2, 4, 8, 16, 32, 64:
1553
+ a = base2ba(n, whitespace)
1554
+ self.assertEqual(a, bitarray())
1555
+ a = urandom(60)
1556
+ c = list(ba2base(n, a))
1557
+ for _ in range(randrange(80)):
1558
+ c.insert(randint(0, len(c)), choice(whitespace))
1559
+ s = ''.join(c)
1560
+ self.assertEqual(base2ba(n, s), a)
1561
+ self.assertRaises(ValueError, base2ba, 8, "3\u00a0") # no-break space
1562
+ self.assertRaises(ValueError, base2ba, 8, b"3\xa0")
1563
+
1564
+ def test_ba2base_group(self):
1565
+ a = bitarray("001 011 100 111", "little")
1566
+ self.assertEqual(ba2base(8, a, 3), "461 7")
1567
+ self.assertEqual(ba2base(8, a, group=2), "46 17")
1568
+ self.assertEqual(ba2base(8, a, sep="_", group=2), "46_17")
1569
+ self.assertEqual(ba2base(8, a, 2, sep="."), "46.17")
1570
+ for n, s, group, sep, res in [
1571
+ (2, '10100', 2, '-', '10-10-0'),
1572
+ (4, '10 11 00 01', 1, "_", "2_3_0_1"),
1573
+ (8, "101 100 011 101 001 010", 3, " ", "543 512"),
1574
+ (8, "101 100 011 101 001 010", 3, "", "543512"),
1575
+ (16, '1011 0001 1101 1010 1111', 4, "+", "b1da+f"),
1576
+ (32, "10110 00111 01101 01111", 2, ", ", "WH, NP"),
1577
+ (64, "101100 011101 101011 111110 101110", 2, ".", "sd.r+.u"),
1578
+ ]:
1579
+ a = bitarray(s, "big")
1580
+ s = ba2base(n, a, group, sep)
1581
+ self.assertIs(type(s), str)
1582
+ self.assertEqual(s, res)
1583
+
1584
+ def test_empty(self):
1585
+ for n in 2, 4, 8, 16, 32, 64:
1586
+ a = base2ba(n, '')
1587
+ self.assertEqual(a, bitarray())
1588
+ self.assertEqual(ba2base(n, a), '')
1589
+
1590
+ def test_invalid_characters(self):
1591
+ for n, s in ((2, '2'), (4, '4'), (8, '8'), (16, 'g'), (32, '8'),
1592
+ (32, '1'), (32, 'a'), (64, '-'), (64, '_')):
1593
+ msg = ("invalid digit found for base%d, "
1594
+ "got '%s' (0x%02x)" % (n, s, ord(s)))
1595
+ self.assertRaisesMessage(ValueError, msg, base2ba, n, s)
1596
+
1597
+ for n in 2, 4, 8, 16, 32, 64:
1598
+ for s in '_', '@', '[', '\u20ac', '\0', b'\0', b'\x80', b'\xff':
1599
+ self.assertRaises(ValueError, base2ba, n, s)
1600
+ msg = "invalid digit found for base%d, got '{' (0x7b)" % n
1601
+ self.assertRaisesMessage(ValueError, msg, base2ba, n, '{')
1602
+
1603
+ def test_invalid_args(self):
1604
+ a = bitarray()
1605
+ self.assertRaises(TypeError, ba2base, None, a)
1606
+ self.assertRaises(TypeError, base2ba, None, '')
1607
+ self.assertRaises(TypeError, ba2base, 16.0, a)
1608
+ self.assertRaises(TypeError, base2ba, 16.0, '')
1609
+ self.assertRaises(TypeError, ba2base, 32, None)
1610
+ self.assertRaises(TypeError, base2ba, 32, None)
1611
+
1612
+ for values, msg in [
1613
+ ([-1023, -16, -1, 0, 3, 5, 31, 48, 63, 129, 511, 4123,
1614
+ 1, 128, 256, 512, 1024, 2048, 4096, 8192],
1615
+ "base must be 2, 4, 8, 16, 32 or 64")]:
1616
+ for i in values:
1617
+ self.assertRaisesMessage(ValueError, msg, ba2base, i, a)
1618
+ self.assertRaisesMessage(ValueError, msg, base2ba, i, '')
1619
+
1620
+ a = bitarray(29)
1621
+ for m in range(2, 7):
1622
+ msg = "bitarray length 29 not multiple of %d" % m
1623
+ self.assertRaisesMessage(ValueError, msg, ba2base, 1 << m, a)
1624
+
1625
+ def test_hexadecimal(self):
1626
+ a = base2ba(16, 'F61', 'big')
1627
+ self.assertEqual(a, bitarray('1111 0110 0001'))
1628
+ self.assertEqual(ba2base(16, a), 'f61')
1629
+
1630
+ for n in range(50):
1631
+ s = ''.join(choices(hexdigits, k=n))
1632
+ endian = choice(ENDIANS)
1633
+ a = base2ba(16, s, endian)
1634
+ self.assertEQUAL(a, hex2ba(s, endian))
1635
+ self.assertEqual(ba2base(16, a), ba2hex(a))
1636
+
1637
+ def test_base32(self):
1638
+ msg = os.urandom(randint(10, 100) * 5)
1639
+ s = base64.b32encode(msg).decode()
1640
+ a = base2ba(32, s, 'big')
1641
+ self.assertEqual(a.tobytes(), msg)
1642
+ self.assertEqual(ba2base(32, a), s)
1643
+ self.assertEqual(base64.b32decode(s), msg)
1644
+
1645
+ def test_base64(self):
1646
+ msg = os.urandom(randint(10, 100) * 3)
1647
+ s = base64.standard_b64encode(msg).decode()
1648
+ a = base2ba(64, s, 'big')
1649
+ self.assertEqual(a.tobytes(), msg)
1650
+ self.assertEqual(ba2base(64, a), s)
1651
+ self.assertEqual(base64.standard_b64decode(s), msg)
1652
+
1653
+ def test_primes(self):
1654
+ primes = gen_primes(60, odd=True)
1655
+ base_2 = primes.to01()
1656
+ for n, endian, rep in [
1657
+ ( 2, "little", base_2),
1658
+ ( 2, "big", base_2),
1659
+ ( 4, "little", "232132030132012122122010132110"),
1660
+ ( 4, "big", "131231030231021211211020231220"),
1661
+ ( 8, "little", "65554155441515405550"),
1662
+ ( 8, "big", "35551455114545105550"),
1663
+ (16, "little", "e6bc4b46a921d61"),
1664
+ (16, "big", "76d32d265948b68"),
1665
+ (32, "little", "O3SJLSJTSI3C"),
1666
+ (32, "big", "O3JS2JSZJC3I"),
1667
+ (64, "little", "utMtkppEtF"),
1668
+ (64, "big", "dtMtJllIto"),
1669
+ ]:
1670
+ a = bitarray(primes, endian)
1671
+ s = ba2base(n, a)
1672
+ self.assertIs(type(s), str)
1673
+ self.assertEqual(s, rep)
1674
+ b = base2ba(n, rep, endian)
1675
+ self.assertEqual(b, a)
1676
+ self.assertIs(type(b), bitarray)
1677
+ self.assertEqual(b.endian, endian)
1678
+
1679
+ alphabets = [
1680
+ # m n alphabet
1681
+ (1, 2, '01'),
1682
+ (2, 4, '0123'),
1683
+ (3, 8, '01234567'),
1684
+ (4, 16, '0123456789abcdef'),
1685
+ (4, 16, '0123456789ABCDEF'),
1686
+ (5, 32, 'ABCDEFGHIJKLMNOPQRSTUVWXYZ234567'),
1687
+ (6, 64, 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdef'
1688
+ 'ghijklmnopqrstuvwxyz0123456789+/'),
1689
+ ]
1690
+
1691
+ def test_alphabets(self):
1692
+ for m, n, alphabet in self.alphabets:
1693
+ self.assertEqual(1 << m, n)
1694
+ self.assertEqual(len(alphabet), n)
1695
+ for i, c in enumerate(alphabet):
1696
+ endian = choice(ENDIANS)
1697
+ self.assertEqual(ba2int(base2ba(n, c, endian)), i)
1698
+ if m == 4 and c in "ABCDEF":
1699
+ c = chr(ord(c) + 32)
1700
+ self.assertEqual(ba2base(n, int2ba(i, m, endian)), c)
1701
+
1702
+ def test_not_alphabets(self):
1703
+ for m, n, alphabet in self.alphabets:
1704
+ for i in range(256):
1705
+ c = chr(i)
1706
+ if c in alphabet or c in whitespace:
1707
+ continue
1708
+ if n == 16 and c in hexdigits:
1709
+ continue
1710
+ self.assertRaises(ValueError, base2ba, n, c)
1711
+
1712
+ def test_random(self):
1713
+ for _ in range(100):
1714
+ m = randint(1, 6)
1715
+ a = urandom_2(m * randrange(100))
1716
+ n = 1 << m
1717
+ s = ba2base(n, a, group=randrange(10), sep=randrange(5) * " ")
1718
+ if m == 4 and getrandbits(1):
1719
+ s = s.upper()
1720
+ if getrandbits(1):
1721
+ s = s.encode()
1722
+ b = base2ba(n, s, a.endian)
1723
+ self.assertEQUAL(a, b)
1724
+ self.check_obj(b)
1725
+
1726
+ # --------------------------- sparse compression ----------------------------
1727
+
1728
+ class SC_Util:
1729
+
1730
+ def get_max_pop(self, n, m=256):
1731
+ """
1732
+ Return the maximal population for which one type `n` block is
1733
+ preferred over `m` blocks of type `n-1`, or 0 if there is none.
1734
+ """
1735
+ self.assertIn(n, (2, 3, 4))
1736
+ self.assertTrue(1 <= m <= 256)
1737
+ h = 1 if n == 2 else 2 # header size for type n-1 block
1738
+ # After cancelling the index bytes common to both alternatives,
1739
+ # `m` type `n-1` blocks cost `h * m` bytes, while one type `n`
1740
+ # block costs `2 + k` bytes.
1741
+ #
1742
+ # h * m = 2 + k -> k = h * m - 2
1743
+ #
1744
+ k = h * m - 2 # population where type n-1 and type n tie
1745
+ # At the tie (k), the encoder prefers type n-1.
1746
+ return max(0, min(255, k - 1))
1747
+
1748
+ def sample_with_highest(self, n, k):
1749
+ "Return `k` samples from `range(n)`, with `n-1` first."
1750
+ self.assertGreater(n, 0)
1751
+ self.assertGreater(k, 0)
1752
+ indices = [n - 1]
1753
+ indices.extend(sample(range(n - 1), k - 1))
1754
+ return indices
1755
+
1756
+ @staticmethod
1757
+ def header(nbits):
1758
+ "little-endian header using 32-bit struct.pack() for nbits"
1759
+ length = bits2bytes(nbits.bit_length())
1760
+ b = bytearray([length])
1761
+ b.extend(struct.pack("<I", nbits)[:length])
1762
+ return bytes(b)
1763
+
1764
+ def make_blob(self, nbits, n, indices):
1765
+ "little-endian streams using 32-bit struct.pack() for nbits"
1766
+ self.assertIn(n, (2, 3, 4))
1767
+ self.assertIn(len(indices), range(256))
1768
+
1769
+ b = bytearray(self.header(nbits))
1770
+ b.append(0xc0 + n)
1771
+ b.append(len(indices))
1772
+ for i in sorted(indices):
1773
+ b.extend(struct.pack("<I", i)[:n])
1774
+ b.append(0) # stop byte
1775
+ return bytes(b)
1776
+
1777
+ @staticmethod
1778
+ def block_counts(n):
1779
+ counts = 5 * [0]
1780
+ counts[n] = 1
1781
+ return counts
1782
+
1783
+ def check_stat(self, a, b, n, check_encode=True):
1784
+ stat = sc_stat(b)
1785
+ self.assertEqual(stat['nbits'], len(a))
1786
+ self.assertEqual(stat['blocks'], self.block_counts(n))
1787
+ self.assertEqual(sc_decode(b), a)
1788
+ if check_encode:
1789
+ self.assertEqual(sc_encode(a), b)
1790
+
1791
+
1792
+ class SC_Tests(unittest.TestCase, Util, SC_Util):
1793
+
1794
+ def test_explicit(self):
1795
+ for b, bits, endian in [
1796
+ (b'\x00\0', '', 'little'),
1797
+ (b'\x01\x03\x01\x03\0', '110', 'little'),
1798
+ (b'\x01\x07\x01\x40\0', '0000001', 'little'),
1799
+ (b'\x11\x07\x01\x02\0', '0000001', 'big'),
1800
+ (b'\x01\x10\x02\xf0\x0f\0', '00001111 11110000', 'little'),
1801
+ (b'\x11\x10\xa1\x0c\0', '00000000 00001000', 'big'),
1802
+ (b'\x11\x09\xa1\x08\0', '00000000 1', 'big'),
1803
+ (b'\x01g\xa4abde\0', 97 * '0' + '110110', 'little'),
1804
+ ]:
1805
+ a = bitarray(bits, endian)
1806
+ self.assertEqual(sc_encode(a), b)
1807
+ self.assertEQUAL(sc_decode(b), a)
1808
+
1809
+ stat = sc_stat(b)
1810
+ self.assertEqual(stat['endian'], endian)
1811
+ self.assertEqual(stat['nbits'], len(a))
1812
+
1813
+ def test_skip_block_data(self):
1814
+ b = [0x01, 197, 0xa3, None, None, None, 0x00]
1815
+ # sc_stat() skips the block data without inspecting its items.
1816
+ self.assertEqual(sc_stat(b), {'endian': 'little',
1817
+ 'nbits': 197,
1818
+ 'blocks': [0, 1, 0, 0, 0]})
1819
+ # sc_decode() must process the block data.
1820
+ self.assertRaises(TypeError, sc_decode, b)
1821
+
1822
+ def test_encode_types(self):
1823
+ for a in bitarray('1', 'big'), frozenbitarray('1', 'big'):
1824
+ b = sc_encode(a)
1825
+ self.assertIs(type(b), bytes)
1826
+ self.assertEqual(b, b'\x11\x01\x01\x80\0')
1827
+
1828
+ for a in None, [], 0, 123, b'', b'\x00', 3.14:
1829
+ self.assertRaises(TypeError, sc_encode, a)
1830
+
1831
+ def test_decode_stat_types(self):
1832
+ blob = b'\x11\x03\x01\x20\0'
1833
+ for b in blob, bytearray(blob), list(blob), array.array('B', blob):
1834
+ a = sc_decode(b)
1835
+ self.assertIs(type(a), bitarray)
1836
+ self.assertEqual(a.endian, 'big')
1837
+ self.assertEqual(a.to01(), '001')
1838
+
1839
+ stat = sc_stat(b)
1840
+ self.assertEqual(stat, {'endian': 'big', 'nbits': 3,
1841
+ 'blocks': [1, 0, 0, 0, 0]})
1842
+ self.assertIs(type(stat), dict)
1843
+ self.assertIs(type(stat['endian']), str)
1844
+ self.assertIs(type(stat['nbits']), int)
1845
+ blocks = stat['blocks']
1846
+ self.assertIs(type(blocks), list)
1847
+ self.assertEqual(len(blocks), 5)
1848
+ self.assertTrue(all(type(n) is int for n in blocks))
1849
+
1850
+ def test_zero_explicit(self):
1851
+ for blob, blocks, nbits in [
1852
+ # empty bitarray
1853
+ (b"\x00\0", [0, 0, 0, 0, 0], 0),
1854
+ (b"\x02\x00\x00\xa0\xc2\x00\xc3\x00\xc4\x00\xa0\0",
1855
+ [0, 2, 1, 1, 1], 0),
1856
+ # zero bitarray with length 8
1857
+ (b"\x01\x08\0", [0, 0, 0, 0, 0], 8),
1858
+ (b"\x03\x08\x00\x00\0", [0, 0, 0, 0, 0], 8),
1859
+ (b"\x01\x08\x01\x00\xa0\xc2\x00\xc3\x00\xc4\x00\xa0\xa0\0",
1860
+ [1, 3, 1, 1, 1], 8),
1861
+ ]:
1862
+ it = iter(blob)
1863
+ a = sc_decode(it)
1864
+ self.assertRaises(StopIteration, next, it)
1865
+ self.assertEqual(a.endian, 'little')
1866
+ self.assertEqual(len(a), nbits)
1867
+ self.assertFalse(a.any())
1868
+
1869
+ it = iter(blob)
1870
+ stat = sc_stat(it)
1871
+ self.assertRaises(StopIteration, next, it)
1872
+ self.assertEqual(stat, {'endian': 'little', 'nbits': nbits,
1873
+ 'blocks': blocks})
1874
+
1875
+ self.assertEqual(sc_encode(bitarray(0, 'little')), b'\x00\0')
1876
+ self.assertEqual(sc_encode(bitarray(8, 'little')), b'\x01\x08\0')
1877
+
1878
+ def test_untouch(self):
1879
+ blob = b"\x01\x07\x01\x73\0XYZ"
1880
+
1881
+ stream = iter(blob)
1882
+ self.assertEqual(sc_decode(stream), bitarray("1100111"))
1883
+ self.assertEqual(next(stream), ord('X'))
1884
+
1885
+ stream = iter(blob)
1886
+ stat = sc_stat(stream)
1887
+ self.assertEqual(stat, {'endian': 'little',
1888
+ 'nbits': 7,
1889
+ 'blocks': [1, 0, 0, 0, 0]})
1890
+ self.assertEqual(next(stream), ord('X'))
1891
+
1892
+ @unittest.skipIf(not HAVE_BINOMIALVARIATE, "Python 3.12+ required")
1893
+ def test_concatenate(self):
1894
+ N = 100
1895
+ arrays = [random_p(randrange(500), 1/8) for _ in range(N)]
1896
+ it = iter(b''.join(sc_encode(a) for a in arrays))
1897
+ self.assertEqual([sc_decode(it) for _ in range(N)], arrays)
1898
+ self.assertRaises(StopIteration, next, it)
1899
+
1900
+ def test_decode_stat_errors(self):
1901
+ for f in sc_decode, sc_stat:
1902
+ a = [17, 3, 1, 32, 0]
1903
+
1904
+ self.assertRaises(TypeError, f) # no argument
1905
+ self.assertRaises(TypeError, f, a, None) # extra argument
1906
+ f(a) # no error
1907
+ for x in 256, -1:
1908
+ a[-1] = x
1909
+ self.assertRaises(ValueError, f, a)
1910
+
1911
+ a = [0x02, None]
1912
+ self.assertRaises(TypeError, f, a)
1913
+ for x in None, 3, 3.2, Ellipsis, 'foo':
1914
+ self.assertRaises(TypeError, f, x)
1915
+
1916
+ def test_decode_header_errors(self):
1917
+ for f in sc_decode, sc_stat:
1918
+ # invalid header
1919
+ for c in 0x20, 0x21, 0x40, 0x80, 0xc0, 0xf0, 0xff:
1920
+ self.assertRaisesMessage(ValueError,
1921
+ "invalid header: 0x%02x" % c,
1922
+ f, [c])
1923
+ # invalid block head
1924
+ for c in 0xc0, 0xc1, 0xc5, 0xff:
1925
+ self.assertRaisesMessage(ValueError,
1926
+ "invalid block head: 0x%02x" % c,
1927
+ f, [0x01, 0x10, c])
1928
+
1929
+ def test_decode_header_overflow(self):
1930
+ self.assertRaisesMessage(
1931
+ OverflowError,
1932
+ "sizeof(Py_ssize_t) = %d: cannot read 9 bytes" % PTRSIZE,
1933
+ sc_decode, b'\x09' + 9 * b'\x00')
1934
+
1935
+ self.assertRaisesMessage(
1936
+ ValueError,
1937
+ "read %d bytes got negative value: -1" % PTRSIZE,
1938
+ sc_decode, [PTRSIZE] + PTRSIZE * [0xff])
1939
+
1940
+ if PTRSIZE == 4:
1941
+ self.assertRaisesMessage(
1942
+ OverflowError,
1943
+ "sizeof(Py_ssize_t) = 4: cannot read 5 bytes",
1944
+ sc_decode, b'\x05' + 5 * b'\x00')
1945
+
1946
+ self.assertRaisesMessage(
1947
+ ValueError,
1948
+ "read 4 bytes got negative value: -2147483648",
1949
+ sc_decode, b'\x04\x00\x00\x00\x80')
1950
+
1951
+ def test_decode_errors(self):
1952
+ # too many raw bytes
1953
+ self.assertRaisesMessage(
1954
+ ValueError, "decode error (raw): 0 + 2 > 1",
1955
+ sc_decode, b"\x01\x05\x02\xff\xff\0")
1956
+ self.assertRaisesMessage(
1957
+ ValueError, "decode error (raw): 32 + 3 > 34",
1958
+ sc_decode, b"\x02\x0f\x01\xa0\x03\xff\xff\xff\0")
1959
+
1960
+ # sparse index too high
1961
+ self.assertRaisesMessage(
1962
+ ValueError, "decode error (n=1): 128 >= 128",
1963
+ sc_decode, b"\x01\x80\xa1\x80\0")
1964
+ self.assertRaisesMessage(
1965
+ ValueError, "decode error (n=2): 512 >= 512",
1966
+ sc_decode, b"\x02\x00\x02\xc2\x01\x00\x02\0")
1967
+ self.assertRaisesMessage(
1968
+ ValueError, "decode error (n=3): 32768 >= 32768",
1969
+ sc_decode, b"\x02\x00\x80\xc3\x01\x00\x80\x00\0")
1970
+
1971
+ msg = {4: "read 4 bytes got negative value: -2147483648",
1972
+ 8: "decode error (n=4): 2147483648 >= 16"}
1973
+ self.assertRaisesMessage(
1974
+ ValueError, msg[PTRSIZE],
1975
+ sc_decode, b"\x01\x10\xc4\x01\x00\x00\x00\x80\0")
1976
+
1977
+ msg = {4: "read 4 bytes got negative value: -1",
1978
+ 8: "decode error (n=4): 4294967295 >= 16"}
1979
+ self.assertRaisesMessage(
1980
+ ValueError, msg[PTRSIZE],
1981
+ sc_decode, b"\x01\x10\xc4\x01\xff\xff\xff\xff\0")
1982
+
1983
+ def test_stop_iteration(self):
1984
+ for stream in [b'', b'\x00', b'\x01', b'\x02\x77',
1985
+ b'\x01\x04\x01', b'\x01\x04\xa1', b'\x01\x04\xa0']:
1986
+ for f in sc_decode, sc_stat:
1987
+ self.assertRaises(StopIteration, f, stream)
1988
+
1989
+ def test_decode_ambiguity(self):
1990
+ for n, blob in [
1991
+ # raw:
1992
+ (0, b'\x11\x03\x01\x20\0'), # this is what sc_encode gives us
1993
+ (0, b'\x11\x03\x01\x3f\0'), # set the pad bits to 1
1994
+ (0, b'\x12\x03\x00\x01\x20\0'), # header has redundant zeros
1995
+ # sparse:
1996
+ (1, b'\x11\x03\xa1\x02\0'),
1997
+ (2, b'\x11\x03\xc2\x01\x02\x00\0'),
1998
+ (3, b'\x11\x03\xc3\x01\x02\x00\x00\0'),
1999
+ (4, b'\x11\x03\xc4\x01\x02\x00\x00\x00\0'),
2000
+ ]:
2001
+ a = sc_decode(blob)
2002
+ self.assertEqual(a.to01(), '001')
2003
+ stat = sc_stat(blob)
2004
+ self.assertEqual(stat['endian'], 'big')
2005
+ self.assertEqual(stat['nbits'], 3)
2006
+ self.assertEqual(stat['blocks'], self.block_counts(n))
2007
+
2008
+ def test_block_type0(self):
2009
+ for k in range(0x01, 0xa0):
2010
+ nbytes = k if k <= 32 else 32 * (k - 31)
2011
+ nbits = 8 * nbytes
2012
+ a = ones(nbits, "little")
2013
+ b = bytearray(self.header(nbits))
2014
+ b.append(k)
2015
+ b.extend(a.tobytes())
2016
+ b.append(0) # stop byte
2017
+
2018
+ self.check_stat(a, b, 0)
2019
+
2020
+ def test_block_type1(self):
2021
+ a = bitarray(256, 'little')
2022
+ indices = sample(range(256), 31)
2023
+
2024
+ for k, index in enumerate(indices, 1):
2025
+ a[index] = 1
2026
+ b = bytearray([0x02, 0x00, 0x01, 0xa0 + k])
2027
+ b.extend(sorted(indices[:k]))
2028
+ b.append(0) # stop byte
2029
+
2030
+ self.check_stat(a, b, 1)
2031
+
2032
+ def test_block_type234(self):
2033
+ # nbits and k are redundant, but make the expected values explicit;
2034
+ # k also tests get_max_pop().
2035
+ for n, m, nbits, k in [
2036
+ (2, 256, 1 << 16, 253),
2037
+ (3, 256, 1 << 24, 255),
2038
+ (4, 4, 1 << 26, 5),
2039
+ ]:
2040
+ prev_size = 1 << ((n - 1) * 8) # size of type n-1 block
2041
+ self.assertEqual(prev_size * m, nbits)
2042
+ a = bitarray(nbits, 'little')
2043
+ self.assertEqual(self.get_max_pop(n, m), k)
2044
+ indices = self.sample_with_highest(nbits, k)
2045
+ a[indices] = 1
2046
+ b = self.make_blob(nbits, n, indices)
2047
+ self.check_stat(a, b, n)
2048
+
2049
+ def test_sparse_population_limit(self):
2050
+ a = zeros(1 << 24)
2051
+ a[::1 << 16] = 1 # 256 populated type-2 blocks
2052
+ b = sc_encode(a)
2053
+ self.assertEqual(sc_decode(b), a)
2054
+ # The index count byte can represent at most 255 indices (0xff), so
2055
+ # the 256 indices do not fit in one type-3 block. The encoder first
2056
+ # uses one type-2 block, leaving 255 indices which fit in one type-3
2057
+ # block.
2058
+ self.assertEqual(sc_stat(b)['blocks'], [0, 0, 1, 1, 0])
2059
+
2060
+ def test_decode_random_bytes(self):
2061
+ # ensure random input doesn't crash the decoder
2062
+ for _ in range(100):
2063
+ n = randrange(20)
2064
+ b = b'\x02\x00\x04' + os.urandom(n)
2065
+ try:
2066
+ a = sc_decode(b)
2067
+ except (StopIteration, ValueError):
2068
+ continue
2069
+ self.assertEqual(len(a), 1024)
2070
+ self.assertEqual(a.endian, 'little')
2071
+
2072
+ def check_blob_length(self, a, m, blocks):
2073
+ blob = sc_encode(a)
2074
+ self.assertEqual(len(blob), m)
2075
+ self.assertEqual(sc_decode(blob), a)
2076
+ stat = sc_stat(blob)
2077
+ self.assertEqual(stat['nbits'], len(a))
2078
+ self.assertEqual(stat['blocks'], blocks)
2079
+
2080
+ def test_encode_zeros(self):
2081
+ for i in range(26):
2082
+ nbits = 1 << i
2083
+ a = zeros(nbits)
2084
+ m = 2 # head byte and stop byte
2085
+ m += bits2bytes(nbits.bit_length()) # size of nbits in bytes
2086
+ self.check_blob_length(a, m, [0, 0, 0, 0, 0])
2087
+
2088
+ a[0] = 1
2089
+ # As the first block accounts for the entire population, trailing
2090
+ # zeros do not cause the block type to increase with array size.
2091
+ n = int(i > 3)
2092
+ m += 2 # block head and raw byte (type 0) or index (type 1)
2093
+
2094
+ self.check_blob_length(a, m, self.block_counts(n))
2095
+
2096
+ def test_encode_trailing_zeros(self):
2097
+ for n, index, blocks in [
2098
+ (1, 0, self.block_counts(1)),
2099
+ (2, 1 << 8, [0, 2, 0, 0, 0]), # 2 type 1 blocks are smaller
2100
+ (2, 3 << 8, self.block_counts(2)),
2101
+ (3, 1 << 16, self.block_counts(3)),
2102
+ ]:
2103
+ # The bitarray is 4 times the decoded size of a type n block.
2104
+ a = bitarray(4 << (8 * n))
2105
+ a[index] = 1
2106
+ blob = sc_encode(a)
2107
+ # Trailing zeros must not widen the encoding.
2108
+ self.assertEqual(sc_stat(blob)['blocks'], blocks)
2109
+ self.assertEqual(sc_decode(blob), a)
2110
+
2111
+ def test_encode_ones(self):
2112
+ for _ in range(10):
2113
+ nbits = randrange(100_000)
2114
+ a = ones(nbits)
2115
+ m = 2 # head byte and stop byte
2116
+ m += bits2bytes(nbits.bit_length()) # size bytes
2117
+ nbytes = bits2bytes(nbits)
2118
+ m += nbytes # actual raw bytes
2119
+ # number of raw blocks:
2120
+ raw_blocks = (
2121
+ bool(nbytes % 32) + # number in 0x01 .. 0x1f
2122
+ (nbytes // 32 + 127) // 128) # number in 0x20 .. 0x9f
2123
+ # number of head bytes
2124
+ m += raw_blocks
2125
+ self.check_blob_length(a, m, [raw_blocks, 0, 0, 0, 0])
2126
+
2127
+ def round_trip(self, a):
2128
+ c = a.copy()
2129
+ blob = sc_encode(a)
2130
+ b = sc_decode(blob)
2131
+ self.assertTrue(a == b == c)
2132
+ self.assertTrue(a.endian == b.endian == c.endian)
2133
+ stat = sc_stat(blob)
2134
+ self.assertEqual(stat['endian'], a.endian)
2135
+ self.assertEqual(stat['nbits'], len(a))
2136
+
2137
+ def test_random(self):
2138
+ for _ in range(5):
2139
+ n = randrange(100_000)
2140
+ endian = choice(ENDIANS)
2141
+ a = ones(n, endian)
2142
+ while a.any():
2143
+ a &= urandom(n, endian)
2144
+ self.round_trip(a)
2145
+
2146
+
2147
+ # ------------------------ variable length codec ------------------------
2148
+
2149
+ class VL_Tests(unittest.TestCase, Util):
2150
+
2151
+ def test_explicit(self):
2152
+ for blob, s in [
2153
+ (b'\x40', ''),
2154
+ (b'\x30', '0'),
2155
+ (b'\x38', '1'),
2156
+ (b'\x00', '0000'),
2157
+ (b'\x01', '0001'),
2158
+ (b'\xd3\x20', '001101'),
2159
+ (b'\xe0\x40', '0000 1'),
2160
+ (b'\x90\x02', '0000 000001'),
2161
+ (b'\xb5\xa7\x18', '0101 0100111 0011'),
2162
+ (b'\x95\xb7\x1c', '0101 0110111 001110'),
2163
+ ]:
2164
+ a = bitarray(s)
2165
+ self.assertEqual(vl_encode(a), blob)
2166
+ c = vl_decode(blob)
2167
+ self.assertEqual(c, a)
2168
+ self.assertEqual(c.endian, get_default_endian())
2169
+
2170
+ for endian in OPT_ENDIANS:
2171
+ a = bitarray(s, endian)
2172
+ c = vl_encode(a)
2173
+ self.assertIs(type(c), bytes)
2174
+ self.assertEqual(c, blob)
2175
+
2176
+ c = vl_decode(blob, endian)
2177
+ self.assertEqual(c, a)
2178
+ self.assertEqual(c.endian, endian or get_default_endian())
2179
+
2180
+ def test_encode_types(self):
2181
+ s = "0011 01"
2182
+ for a in bitarray(s), frozenbitarray(s):
2183
+ b = vl_encode(a)
2184
+ self.assertIs(type(b), bytes)
2185
+ self.assertEqual(b, b'\xd3\x20')
2186
+
2187
+ for a in None, [], 0, 123, b'', b'\x00', 3.14:
2188
+ self.assertRaises(TypeError, vl_encode, a)
2189
+
2190
+ def test_decode_types(self):
2191
+ blob = b'\xd3\x20'
2192
+ for s in (blob, iter(blob), memoryview(blob), iter([0xd3, 0x20]),
2193
+ bytearray(blob)):
2194
+ a = vl_decode(s, endian=choice(ENDIANS))
2195
+ self.assertIs(type(a), bitarray)
2196
+ self.assertEqual(a, bitarray('0011 01'))
2197
+
2198
+ # these objects are not iterable
2199
+ for arg in None, 0, 1, 0.0:
2200
+ self.assertRaises(TypeError, vl_decode, arg)
2201
+ # these items cannot be interpreted as ints
2202
+ for item in None, 2.34, Ellipsis, 'foo':
2203
+ self.assertRaises(TypeError, vl_decode, iter([0x95, item]))
2204
+
2205
+ def test_decode_args(self):
2206
+ # item not integer
2207
+ self.assertRaises(TypeError, vl_decode, iter([b'\x40']))
2208
+
2209
+ self.assertRaises(TypeError, vl_decode, b'\x40', 'big', 3)
2210
+ self.assertRaises(ValueError, vl_decode, b'\x40', 'foo')
2211
+
2212
+ def test_decode_trailing(self):
2213
+ for s, bits in [(b'\x40ABC', ''),
2214
+ (b'\xe0\x40A', '00001')]:
2215
+ stream = iter(s)
2216
+ self.assertEqual(vl_decode(stream), bitarray(bits))
2217
+ self.assertEqual(next(stream), 65)
2218
+
2219
+ def test_decode_ambiguity(self):
2220
+ for s in b'\x40', b'\x4f', b'\x45':
2221
+ self.assertEqual(vl_decode(s), bitarray())
2222
+ for s in b'\x1e', b'\x1f':
2223
+ self.assertEqual(vl_decode(s), bitarray('111'))
2224
+
2225
+ def test_decode_stream(self):
2226
+ stream = iter(b'\x40\x30\x38\x40\x2c\xe0\x40\xd3\x20')
2227
+ for bits in '', '0', '1', '', '11', '0000 1', '0011 01':
2228
+ self.assertEqual(vl_decode(stream), bitarray(bits))
2229
+
2230
+ arrays = [urandom(randrange(30)) for _ in range(1000)]
2231
+ stream = iter(b''.join(vl_encode(a) for a in arrays))
2232
+ for a in arrays:
2233
+ self.assertEqual(vl_decode(stream), a)
2234
+
2235
+ def test_decode_errors(self):
2236
+ # decode empty bytes
2237
+ self.assertRaises(StopIteration, vl_decode, b'')
2238
+ # invalid head byte
2239
+ for s in [
2240
+ b'\x70', b'\xf0', # padding = 7
2241
+ b'\x50', b'\x60', b'\x70', # no second byte, but padding > 4
2242
+ ]:
2243
+ self.assertRaisesMessage(ValueError,
2244
+ "invalid head byte: 0x%02x" % s[0],
2245
+ vl_decode, s)
2246
+ # high bit set, but no terminating byte
2247
+ for s in b'\x80', b'\x80\x80':
2248
+ self.assertRaises(StopIteration, vl_decode, s)
2249
+ # decode list with out of range items
2250
+ for i in -1, 256:
2251
+ self.assertRaises(ValueError, vl_decode, [i])
2252
+ # wrong type
2253
+ self.assertRaises(TypeError, vl_decode, [None])
2254
+
2255
+ def test_decode_invalid_stream(self):
2256
+ N = 100
2257
+ s = iter(N * (3 * [0x80] + ['XX']) + ['end.'])
2258
+ for _ in range(N):
2259
+ a = None
2260
+ try:
2261
+ a = vl_decode(s)
2262
+ except TypeError:
2263
+ pass
2264
+ self.assertIsNone(a)
2265
+ self.assertEqual(next(s), 'end.')
2266
+
2267
+ def test_explicit_zeros(self):
2268
+ for n in range(100):
2269
+ a = zeros(4 + n * 7)
2270
+ s = n * b'\x80' + b'\x00'
2271
+ self.assertEqual(vl_encode(a), s)
2272
+ self.assertEqual(vl_decode(s), a)
2273
+
2274
+ @unittest.skipIf(PTRSIZE != 8, 'requires 64 bit machine')
2275
+ def test_uleb128(self):
2276
+ LEN_PAD_BITS = 3
2277
+ for n in range(1, 61):
2278
+ a = urandom(n, 'little')
2279
+ a[-1] = 1 # avoid overlong ULEB128 representation
2280
+ b = vl_encode(a)
2281
+
2282
+ # Treat the VL bytes as canonical ULEB128.
2283
+ payload = bitarray(b, endian='little')
2284
+ del payload[7::8] # remove each byte's continuation bit
2285
+ i = ba2int(payload)
2286
+ self.assertEqual(uleb128_encode(i), b)
2287
+
2288
+ # Remove continuation bits, padding metadata, and actual padding
2289
+ # to recover the original bitarray.
2290
+ c = bitarray(b, 'big')
2291
+ del c[::8] # remove each byte's continuation bit
2292
+ padding = ba2int(c[:LEN_PAD_BITS])
2293
+ self.assertEqual(LEN_PAD_BITS + n + padding, 7 * len(b))
2294
+ del c[:LEN_PAD_BITS] # remove padding metadata
2295
+ del c[len(c) - padding:] # remove actual padding
2296
+ self.assertEqual(c, a)
2297
+
2298
+ def test_concatenate(self):
2299
+ N = 100
2300
+ arrays = [urandom(randrange(20)) for _ in range(N)]
2301
+ it = iter(b''.join(vl_encode(a) for a in arrays))
2302
+ self.assertEqual([vl_decode(it) for _ in range(N)], arrays)
2303
+ self.assertRaises(StopIteration, next, it)
2304
+
2305
+ def round_trip(self, a):
2306
+ c = a.copy()
2307
+ s = vl_encode(a)
2308
+ b = vl_decode(s)
2309
+ self.check_obj(b)
2310
+ self.assertTrue(a == b == c)
2311
+ LEN_PAD_BITS = 3
2312
+ HEAD_WIDTH = 7 - LEN_PAD_BITS # number of payload bits in head
2313
+ self.assertEqual(len(s), (len(a) + LEN_PAD_BITS + 6) // 7)
2314
+ head = s[0]
2315
+ padding = (head & 0x70) >> HEAD_WIDTH
2316
+ self.assertEqual(LEN_PAD_BITS + len(a) + padding, 7 * len(s))
2317
+
2318
+ # overflow-safe C implementation in _util.c
2319
+ r = len(a) % 7
2320
+ self.assertEqual(len(a) // 7 + 1 + (r > HEAD_WIDTH), len(s))
2321
+ self.assertEqual((HEAD_WIDTH - r + 7) % 7, padding)
2322
+
2323
+ def test_large(self):
2324
+ for _ in range(5):
2325
+ a = urandom(randrange(100_000))
2326
+ self.round_trip(a)
2327
+
2328
+ def test_random(self):
2329
+ for n in range(50):
2330
+ a = urandom_2(n)
2331
+ self.round_trip(a)
2332
+
2333
+ # ---------------------------------------------------------------------------
2334
+
2335
+ class IntegerizationTests(unittest.TestCase, Util):
2336
+
2337
+ def test_ba2int(self):
2338
+ self.assertEqual(ba2int(bitarray('0')), 0)
2339
+ self.assertEqual(ba2int(bitarray('1')), 1)
2340
+ self.assertEqual(ba2int(bitarray('00101', 'big')), 5)
2341
+ self.assertEqual(ba2int(bitarray('00101', 'little')), 20)
2342
+ self.assertEqual(ba2int(frozenbitarray('11')), 3)
2343
+ self.assertRaises(ValueError, ba2int, bitarray())
2344
+ self.assertRaises(ValueError, ba2int, frozenbitarray())
2345
+ self.assertRaises(TypeError, ba2int, '101')
2346
+ a = bitarray('111')
2347
+ b = a.copy()
2348
+ self.assertEqual(ba2int(a), 7)
2349
+ # ensure original object wasn't altered
2350
+ self.assertEQUAL(a, b)
2351
+
2352
+ def test_ba2int_frozen(self):
2353
+ for a in self.randombitarrays(start=1):
2354
+ b = frozenbitarray(a)
2355
+ self.assertEqual(ba2int(b), ba2int(a))
2356
+ self.assertEQUAL(a, b)
2357
+
2358
+ def test_ba2int_random(self):
2359
+ for a in self.randombitarrays(start=1):
2360
+ b = bitarray(a, 'big')
2361
+ self.assertEqual(a, b)
2362
+ self.assertEqual(ba2int(b), int(b.to01(), 2))
2363
+
2364
+ def test_ba2int_bytes(self):
2365
+ for n in range(1, 50):
2366
+ a = urandom_2(8 * n)
2367
+ c = bytearray(a.tobytes())
2368
+ i = 0
2369
+ for x in (c if a.endian == 'big' else reversed(c)):
2370
+ i <<= 8
2371
+ i |= x
2372
+ self.assertEqual(ba2int(a), i)
2373
+
2374
+ def test_int2ba(self):
2375
+ self.assertEqual(int2ba(0), bitarray('0'))
2376
+ self.assertEqual(int2ba(1), bitarray('1'))
2377
+ self.assertEqual(int2ba(5), bitarray('101'))
2378
+ self.assertEQUAL(int2ba(6, endian='big'), bitarray('110', 'big'))
2379
+ self.assertEQUAL(int2ba(6, endian='little'),
2380
+ bitarray('011', 'little'))
2381
+
2382
+ def test_errors(self):
2383
+ self.assertRaises(TypeError, int2ba, 1.0)
2384
+ self.assertRaises(TypeError, int2ba, 1, 3.0)
2385
+ self.assertRaises(ValueError, int2ba, 1, 0)
2386
+ self.assertRaises(TypeError, int2ba, 1, 10, 123)
2387
+ self.assertRaises(ValueError, int2ba, 1, 10, 'asd')
2388
+ # signed integer requires length
2389
+ self.assertRaises(TypeError, int2ba, 100, signed=True)
2390
+ # argument cannot be float
2391
+ self.assertRaises(TypeError, int2ba, 123.0)
2392
+
2393
+ def test_signed(self):
2394
+ for s, i in [
2395
+ ('0', 0),
2396
+ ('1', -1),
2397
+ ('00', 0),
2398
+ ('10', 1),
2399
+ ('01', -2),
2400
+ ('11', -1),
2401
+ ('000', 0),
2402
+ ('100', 1),
2403
+ ('010', 2),
2404
+ ('110', 3),
2405
+ ('001', -4),
2406
+ ('101', -3),
2407
+ ('011', -2),
2408
+ ('111', -1),
2409
+ ('00000', 0),
2410
+ ('11110', 15),
2411
+ ('00001', -16),
2412
+ ('11111', -1),
2413
+ ('00000000 0', 0),
2414
+ ('11111111 0', 255),
2415
+ ('00000000 1', -256),
2416
+ ('11111111 1', -1),
2417
+ ]:
2418
+ self.assertEqual(ba2int(bitarray(s, 'little'), signed=1), i)
2419
+ self.assertEqual(ba2int(bitarray(s[::-1], 'big'), signed=1), i)
2420
+
2421
+ len_s = len(bitarray(s))
2422
+ self.assertEQUAL(int2ba(i, len_s, 'little', signed=1),
2423
+ bitarray(s, 'little'))
2424
+ self.assertEQUAL(int2ba(i, len_s, 'big', signed=1),
2425
+ bitarray(s[::-1], 'big'))
2426
+
2427
+ def test_zero(self):
2428
+ for endian in "little", "big":
2429
+ a = int2ba(0, endian=endian)
2430
+ self.assertEQUAL(a, bitarray('0', endian=endian))
2431
+ for n in range(1, 100):
2432
+ a = int2ba(0, length=n, endian=endian, signed=True)
2433
+ b = bitarray(n * '0', endian)
2434
+ self.assertEQUAL(a, b)
2435
+ for signed in 0, 1:
2436
+ self.assertEqual(ba2int(b, signed=signed), 0)
2437
+
2438
+ def test_negative_one(self):
2439
+ for endian in "little", "big":
2440
+ for n in range(1, 100):
2441
+ a = int2ba(-1, length=n, endian=endian, signed=True)
2442
+ b = bitarray(n * '1', endian)
2443
+ self.assertEQUAL(a, b)
2444
+ self.assertEqual(ba2int(b, signed=True), -1)
2445
+
2446
+ def test_int2ba_overflow(self):
2447
+ self.assertRaises(OverflowError, int2ba, -1)
2448
+ self.assertRaises(OverflowError, int2ba, -1, 4)
2449
+
2450
+ self.assertRaises(OverflowError, int2ba, 128, 7)
2451
+ self.assertRaises(OverflowError, int2ba, 64, 7, signed=1)
2452
+ self.assertRaises(OverflowError, int2ba, -65, 7, signed=1)
2453
+
2454
+ for n in range(1, 20):
2455
+ self.assertRaises(OverflowError, int2ba, 1 << n, n)
2456
+ self.assertRaises(OverflowError, int2ba, 1 << (n - 1), n,
2457
+ signed=1)
2458
+ self.assertRaises(OverflowError, int2ba, -(1 << (n - 1)) - 1, n,
2459
+ signed=1)
2460
+
2461
+ def test_int2ba_length(self):
2462
+ self.assertRaises(TypeError, int2ba, 0, 1.0)
2463
+ self.assertRaises(ValueError, int2ba, 0, 0)
2464
+ self.assertEqual(int2ba(5, length=6, endian='big'),
2465
+ bitarray('000101'))
2466
+ for n in range(1, 100):
2467
+ ab = int2ba(1, n, 'big')
2468
+ al = int2ba(1, n, 'little')
2469
+ self.assertEqual(ab.endian, 'big')
2470
+ self.assertEqual(al.endian, 'little')
2471
+ self.assertEqual(len(ab), n),
2472
+ self.assertEqual(len(al), n)
2473
+ self.assertEqual(ab, bitarray((n - 1) * '0') + bitarray('1'))
2474
+ self.assertEqual(al, bitarray('1') + bitarray((n - 1) * '0'))
2475
+
2476
+ ab = int2ba(0, n, 'big')
2477
+ al = int2ba(0, n, 'little')
2478
+ self.assertEqual(len(ab), n)
2479
+ self.assertEqual(len(al), n)
2480
+ self.assertEqual(ab, bitarray(n * '0', 'big'))
2481
+ self.assertEqual(al, bitarray(n * '0', 'little'))
2482
+
2483
+ self.assertEqual(int2ba(2 ** n - 1), bitarray(n * '1'))
2484
+ self.assertEqual(int2ba(2 ** n - 1, endian='little'),
2485
+ bitarray(n * '1'))
2486
+
2487
+ def test_explicit(self):
2488
+ for i, sa in [( 0, '0'), (1, '1'),
2489
+ ( 2, '10'), (3, '11'),
2490
+ (25, '11001'), (265, '100001001'),
2491
+ (3691038, '1110000101001000011110')]:
2492
+ ab = bitarray(sa, 'big')
2493
+ al = bitarray(sa[::-1], 'little')
2494
+ self.assertEQUAL(int2ba(i),
2495
+ ab if get_default_endian() == 'big' else al)
2496
+ self.assertEQUAL(int2ba(i, endian='big'), ab)
2497
+ self.assertEQUAL(int2ba(i, endian='little'), al)
2498
+ self.assertEqual(ba2int(ab), ba2int(al), i)
2499
+
2500
+ def check_round_trip(self, i):
2501
+ for endian in 'big', 'little':
2502
+ a = int2ba(i, endian=endian)
2503
+ self.check_obj(a)
2504
+ self.assertEqual(a.endian, endian)
2505
+ self.assertTrue(len(a) > 0)
2506
+ # ensure we have no leading zeros
2507
+ if a.endian == 'big':
2508
+ self.assertTrue(len(a) == 1 or a.index(1) == 0)
2509
+ self.assertEqual(ba2int(a), i)
2510
+ if i > 0:
2511
+ self.assertEqual(i.bit_length(), len(a))
2512
+ # add a few trailing / leading zeros to bitarray
2513
+ if endian == 'big':
2514
+ a = zeros(randrange(4), endian) + a
2515
+ else:
2516
+ a = a + zeros(randrange(4), endian)
2517
+ self.assertEqual(a.endian, endian)
2518
+ self.assertEqual(ba2int(a), i)
2519
+
2520
+ def test_many(self):
2521
+ for _ in range(20):
2522
+ self.check_round_trip(randrange(10 ** randint(3, 300)))
2523
+
2524
+ @staticmethod
2525
+ def twos_complement(i, num_bits):
2526
+ # https://en.wikipedia.org/wiki/Two%27s_complement
2527
+ mask = 2 ** (num_bits - 1)
2528
+ return -(i & mask) + (i & ~mask)
2529
+
2530
+ def test_random_signed(self):
2531
+ for a in self.randombitarrays(start=1):
2532
+ i = ba2int(a, signed=True)
2533
+ b = int2ba(i, len(a), a.endian, signed=True)
2534
+ self.assertEQUAL(a, b)
2535
+
2536
+ j = ba2int(a, signed=False) # unsigned
2537
+ if i >= 0:
2538
+ self.assertEqual(i, j)
2539
+
2540
+ self.assertEqual(i, self.twos_complement(j, len(a)))
2541
+
2542
+ # ---------------------------------------------------------------------------
2543
+
2544
+ class MixedTests(unittest.TestCase, Util):
2545
+
2546
+ def test_bin(self):
2547
+ for _ in range(20):
2548
+ i = randrange(1000)
2549
+ s = bin(i)
2550
+ self.assertEqual(s[:2], '0b')
2551
+ a = bitarray(s[2:], 'big')
2552
+ self.assertEqual(ba2int(a), i)
2553
+ t = a.to01()
2554
+ self.assertEqual(t, s[2:])
2555
+ self.assertEqual(int(t, 2), i)
2556
+
2557
+ def test_oct(self):
2558
+ for _ in range(20):
2559
+ i = randrange(1000)
2560
+ s = oct(i)
2561
+ self.assertEqual(s[:2], '0o')
2562
+ a = base2ba(8, s[2:], 'big')
2563
+ self.assertEqual(ba2int(a), i)
2564
+ t = ba2base(8, a)
2565
+ self.assertEqual(t, s[2:])
2566
+ self.assertEqual(int(t, 8), i)
2567
+
2568
+ def test_hex(self):
2569
+ for _ in range(20):
2570
+ i = randrange(1000)
2571
+ s = hex(i)
2572
+ self.assertEqual(s[:2], '0x')
2573
+ a = hex2ba(s[2:], 'big')
2574
+ self.assertEqual(ba2int(a), i)
2575
+ t = ba2hex(a)
2576
+ self.assertEqual(t, s[2:])
2577
+ self.assertEqual(int(t, 16), i)
2578
+
2579
+ def test_bitwise(self):
2580
+ for a in self.randombitarrays(start=1):
2581
+ b = urandom(len(a), a.endian)
2582
+ aa = a.copy()
2583
+ bb = b.copy()
2584
+ i = ba2int(a)
2585
+ j = ba2int(b)
2586
+ self.assertEqual(ba2int(a & b), i & j)
2587
+ self.assertEqual(ba2int(a | b), i | j)
2588
+ self.assertEqual(ba2int(a ^ b), i ^ j)
2589
+
2590
+ n = randint(0, len(a))
2591
+ if a.endian == 'big':
2592
+ self.assertEqual(ba2int(a >> n), i >> n)
2593
+ c = zeros(len(a), 'big') + a
2594
+ self.assertEqual(ba2int(c << n), i << n)
2595
+
2596
+ self.assertEQUAL(a, aa)
2597
+ self.assertEQUAL(b, bb)
2598
+
2599
+ def test_bitwise_inplace(self):
2600
+ for a in self.randombitarrays(start=1):
2601
+ b = urandom(len(a), a.endian)
2602
+ bb = b.copy()
2603
+ i = ba2int(a)
2604
+ j = ba2int(b)
2605
+ c = a.copy()
2606
+ c &= b
2607
+ self.assertEqual(ba2int(c), i & j)
2608
+ c = a.copy()
2609
+ c |= b
2610
+ self.assertEqual(ba2int(c), i | j)
2611
+ c = a.copy()
2612
+ c ^= b
2613
+ self.assertEqual(ba2int(c), i ^ j)
2614
+ self.assertEQUAL(b, bb)
2615
+
2616
+ n = randint(0, len(a))
2617
+ if a.endian == 'big':
2618
+ c = a.copy()
2619
+ c >>= n
2620
+ self.assertEqual(ba2int(c), i >> n)
2621
+ c = zeros(len(a), 'big') + a
2622
+ c <<= n
2623
+ self.assertEqual(ba2int(c), i << n)
2624
+
2625
+ # ---------------------- serialize() deserialize() -----------------------
2626
+
2627
+ class SerializationTests(unittest.TestCase, Util):
2628
+
2629
+ def test_explicit(self):
2630
+ for blob, endian, bits in [
2631
+ (b'\x00', 'little', ''),
2632
+ (b'\x07\x01', 'little', '1'),
2633
+ (b'\x17\x80', 'big', '1'),
2634
+ (b'\x13\xf8', 'big', '11111'),
2635
+ (b'\x00\x0f', 'little', '11110000'),
2636
+ (b'\x10\xf0', 'big', '11110000'),
2637
+ (b'\x12\x87\xd8', 'big', '10000111 110110')
2638
+ ]:
2639
+ a = bitarray(bits, endian)
2640
+ s = serialize(a)
2641
+ self.assertEqual(blob, s)
2642
+ self.assertIs(type(s), bytes)
2643
+
2644
+ b = deserialize(blob)
2645
+ self.assertEqual(b, a)
2646
+ self.assertEqual(b.endian, endian)
2647
+ self.assertIs(type(b), bitarray)
2648
+
2649
+ def test_serialize_args(self):
2650
+ for x in '0', 0, 1, b'\x00', 0.0, [0, 1], bytearray([0]):
2651
+ self.assertRaises(TypeError, serialize, x)
2652
+ # no arguments
2653
+ self.assertRaises(TypeError, serialize)
2654
+ # too many arguments
2655
+ self.assertRaises(TypeError, serialize, bitarray(), 1)
2656
+
2657
+ for a in bitarray('0111', 'big'), frozenbitarray('0111', 'big'):
2658
+ self.assertEqual(serialize(a), b'\x14\x70')
2659
+
2660
+ def test_deserialize_args(self):
2661
+ for x in 0, 1, False, True, None, '', '01', 0.0, [0, 1]:
2662
+ self.assertRaises(TypeError, deserialize, x)
2663
+ # no arguments
2664
+ self.assertRaises(TypeError, deserialize)
2665
+ # too many arguments
2666
+ self.assertRaises(TypeError, deserialize, b'\x00', 1)
2667
+
2668
+ blob = b'\x03\x06'
2669
+ x = bitarray(blob)
2670
+ for s in blob, bytearray(blob), memoryview(blob), x:
2671
+ a = deserialize(s)
2672
+ self.assertEqual(a.to01(), '01100')
2673
+ self.assertEqual(a.endian, 'little')
2674
+
2675
+ def test_invalid_bytes(self):
2676
+ self.assertRaises(ValueError, deserialize, b'')
2677
+
2678
+ def check_msg(b):
2679
+ msg = "invalid header byte: 0x%02x" % b[0]
2680
+ self.assertRaisesMessage(ValueError, msg, deserialize, b)
2681
+
2682
+ for i in range(256):
2683
+ b = bytearray([i])
2684
+ if i == 0 or i == 16:
2685
+ self.assertEqual(deserialize(b), bitarray())
2686
+ else:
2687
+ self.assertRaises(ValueError, deserialize, b)
2688
+ check_msg(b)
2689
+
2690
+ b.append(0)
2691
+ if i < 32 and i % 16 < 8:
2692
+ self.assertEqual(deserialize(b), zeros(8 - i % 8))
2693
+ else:
2694
+ self.assertRaises(ValueError, deserialize, b)
2695
+ check_msg(b)
2696
+
2697
+ def test_padbits_ignored(self):
2698
+ for blob, endian in [
2699
+ (b'\x07\x01', 'little'),
2700
+ (b'\x07\x03', 'little'),
2701
+ (b'\x07\xff', 'little'),
2702
+ (b'\x17\x80', 'big'),
2703
+ (b'\x17\xc0', 'big'),
2704
+ (b'\x17\xff', 'big'),
2705
+ ]:
2706
+ a = deserialize(blob)
2707
+ self.assertEqual(a.to01(), '1')
2708
+ self.assertEqual(a.endian, endian)
2709
+
2710
+ def test_random(self):
2711
+ for a in self.randombitarrays():
2712
+ b = serialize(a)
2713
+ c = deserialize(b)
2714
+ self.assertEqual(a, c)
2715
+ self.assertEqual(a.endian, c.endian)
2716
+ self.check_obj(c)
2717
+
2718
+ # ------------------------------ Huffman coding ---------------------------
2719
+
2720
+ class HuffmanUtil:
2721
+
2722
+ def check_code(self, code):
2723
+ self.assertIs(type(code), dict)
2724
+ for v in code.values():
2725
+ self.assertIs(type(v), frozenbitarray)
2726
+ self.assertGreater(len(v), 0)
2727
+
2728
+ if sys.version_info[:2] >= (3, 15):
2729
+ hash(frozendict(code))
2730
+
2731
+ tree = decodetree(code)
2732
+ self.assertEqual(tree.todict(), code)
2733
+ n = len(code)
2734
+ nodes = tree.nodes()
2735
+ if n == 1:
2736
+ # one-symbol code is incomplete
2737
+ self.assertEqual(nodes, (1, 0, 1))
2738
+ else:
2739
+ # proper Huffman tree is complete
2740
+ self.assertEqual(nodes, (0, n - 1, n))
2741
+
2742
+
2743
+ class CommonHuffmanTests(HuffmanUtil, unittest.TestCase):
2744
+
2745
+ # tests for both huffman_code() and canonical_huffman()
2746
+
2747
+ def test_wrong_arg(self):
2748
+ for F in huffman_code, canonical_huffman:
2749
+ self.assertRaises(TypeError, F)
2750
+ self.assertRaises(TypeError, F, [('a', 1)])
2751
+ self.assertRaises(TypeError, F, 123)
2752
+ self.assertRaises(TypeError, F, None)
2753
+ # cannot compare 'a' with 1
2754
+ self.assertRaises(TypeError, F, {'A': 'a', 'B': 1})
2755
+ # frequency map cannot be empty
2756
+ d = {}
2757
+ self.assertRaises(ValueError, F, d)
2758
+ self.assertEqual(d, {}) # d wasn't changed by F
2759
+ # wrong endianness type / extra argument
2760
+ self.assertRaises(TypeError, F, {'A': 2, 'B': 1}, 1)
2761
+
2762
+ @staticmethod
2763
+ def create_codes(freq):
2764
+ yield huffman_code(freq)
2765
+ yield canonical_huffman(freq)[0]
2766
+
2767
+ def test_small_range(self):
2768
+ for n in range(1, 10):
2769
+ f = {i: randrange(10) for i in range(n)}
2770
+ f2 = dict(f)
2771
+ for code in self.create_codes(f):
2772
+ self.assertEqual(len(code), n)
2773
+ self.check_code(code)
2774
+ # make sure frequency map wasn't changed
2775
+ self.assertEqual(f, f2)
2776
+
2777
+ @unittest.skipIf(sys.version_info[:2] < (3, 15),
2778
+ "frozendict introduced in Python 3.15")
2779
+ def test_from_frozendict(self):
2780
+ for n in range(1, 10):
2781
+ f = frozendict({i: random() for i in range(n)})
2782
+ for code in self.create_codes(f):
2783
+ self.assertEqual(len(code), n)
2784
+ self.check_code(code)
2785
+
2786
+ def test_from_counter(self):
2787
+ c = Counter("abracadabra")
2788
+ for code in self.create_codes(c):
2789
+ self.check_code(code)
2790
+
2791
+ def test_large(self):
2792
+ N = randrange(1000, 2000)
2793
+ f = {i: random() for i in range(N)}
2794
+ for code in self.create_codes(f):
2795
+ self.assertEqual(len(code), N)
2796
+ self.check_code(code)
2797
+
2798
+ def test_balanced(self):
2799
+ n = 7
2800
+ freq = {}
2801
+ for i in range(1 << n):
2802
+ freq[i] = 1
2803
+ for code in self.create_codes(freq):
2804
+ self.assertEqual(len(code), 1 << n)
2805
+ for v in code.values():
2806
+ self.assertEqual(len(v), n)
2807
+ self.check_code(code)
2808
+
2809
+ def test_unbalanced(self):
2810
+ n = 37
2811
+ freq = {}
2812
+ for i in range(n):
2813
+ freq[i] = 1 << i
2814
+ for code in self.create_codes(freq):
2815
+ self.assertEqual(len(code), n)
2816
+ for i in range(n):
2817
+ self.assertEqual(len(code[i]), n - max(1, i))
2818
+ self.check_code(code)
2819
+
2820
+ def test_random_list(self):
2821
+ plain = choices(range(100), k=500)
2822
+ for code in self.create_codes(Counter(plain)):
2823
+ a = bitarray()
2824
+ a.encode(code, plain)
2825
+ self.assertEqual(list(a.decode(code)), plain)
2826
+ self.check_code(code)
2827
+
2828
+
2829
+ class HuffmanTests(HuffmanUtil, unittest.TestCase):
2830
+
2831
+ def test_simple(self):
2832
+ freq = {0: 10, 'as': 2, None: 1.6}
2833
+ code = huffman_code(freq)
2834
+ self.assertEqual(len(code), 3)
2835
+ self.assertEqual(len(code[0]), 1)
2836
+ self.assertEqual(len(code['as']), 2)
2837
+ self.assertEqual(len(code[None]), 2)
2838
+
2839
+ def test_endianness_and_type(self):
2840
+ freq = {'A': 10, 'B': 2, 'C': 5}
2841
+ for endian in 'big', 'little':
2842
+ code = huffman_code(freq, endian)
2843
+ self.assertEqual(len(code), 3)
2844
+ for v in code.values():
2845
+ self.assertEqual(v.endian, endian)
2846
+ self.assertIs(type(v), frozenbitarray)
2847
+
2848
+ def test_one_symbol(self):
2849
+ cnt = {'a': 1}
2850
+ endian = choice(ENDIANS)
2851
+ code = huffman_code(cnt, endian)
2852
+ self.assertEqual(code, {'a': bitarray('0')})
2853
+ v = code['a']
2854
+ self.assertEqual(v.endian, endian)
2855
+ self.assertIs(type(v), frozenbitarray)
2856
+ for n in range(4):
2857
+ msg = n * ['a']
2858
+ a = bitarray()
2859
+ a.encode(code, msg)
2860
+ self.assertEqual(a, zeros(n))
2861
+ self.assertEqual(list(a.decode(code)), msg)
2862
+ a.append(1)
2863
+ self.assertRaises(ValueError, list, a.decode(code))
2864
+
2865
+ def test_counter(self):
2866
+ message = 'the quick brown fox jumps over the lazy dog.'
2867
+ code = huffman_code(Counter(message))
2868
+ a = bitarray()
2869
+ a.encode(code, message)
2870
+ self.assertEqual(''.join(a.decode(code)), message)
2871
+ self.check_code(code)
2872
+
2873
+
2874
+ class CanonicalHuffmanTests(Util, HuffmanUtil, unittest.TestCase):
2875
+
2876
+ def test_basic(self):
2877
+ plain = bytearray(b'the quick brown fox jumps over the lazy dog.')
2878
+ chc, count, symbol = canonical_huffman(Counter(plain))
2879
+ self.assertIs(type(chc), dict)
2880
+ for v in chc.values():
2881
+ self.assertIs(type(v), frozenbitarray)
2882
+ self.assertIs(type(count), list)
2883
+ self.assertIs(type(symbol), list)
2884
+ a = bitarray()
2885
+ a.encode(chc, plain)
2886
+ self.assertEqual(bytearray(a.decode(chc)), plain)
2887
+ self.assertEqual(bytearray(canonical_decode(a, count, symbol)), plain)
2888
+
2889
+ def test_example(self):
2890
+ F = frozenbitarray
2891
+ cnt = {'a': 5, 'b': 3, 'c': 1, 'd': 1, 'r': 2}
2892
+ codedict, count, symbol = canonical_huffman(cnt)
2893
+ self.assertEqual(codedict, {'a': F('0'),
2894
+ 'b': F('10'),
2895
+ 'c': F('1110'),
2896
+ 'd': F('1111'),
2897
+ 'r': F('110')})
2898
+ self.assertEqual(count, [0, 1, 1, 1, 2])
2899
+ self.assertEqual(symbol, ['a', 'b', 'r', 'c', 'd'])
2900
+ a = bitarray('01011001110011110101100')
2901
+ msg = "abracadabra"
2902
+ self.assertEqual(''.join(a.decode(codedict)), msg)
2903
+ self.assertEqual(''.join(canonical_decode(a, count, symbol)), msg)
2904
+
2905
+ def test_one_symbol(self):
2906
+ cnt = {'a': 1}
2907
+ chc, count, symbol = canonical_huffman(cnt)
2908
+ self.assertEqual(chc, {'a': bitarray('0')})
2909
+ self.assertIs(type(chc['a']), frozenbitarray)
2910
+ self.assertEqual(count, [0, 1])
2911
+ self.assertEqual(symbol, ['a'])
2912
+ for n in range(4):
2913
+ msg = n * ['a']
2914
+ a = bitarray()
2915
+ a.encode(chc, msg)
2916
+ self.assertEqual(a.to01(), n * '0')
2917
+ self.assertEqual(list(canonical_decode(a, count, symbol)), msg)
2918
+ a.append(1)
2919
+ self.assertRaises(ValueError, list,
2920
+ canonical_decode(a, count, symbol))
2921
+
2922
+ def test_canonical_decode_errors(self):
2923
+ C = canonical_decode
2924
+ a = bitarray('1101')
2925
+ s = ['a']
2926
+ # bitarray not of bitarray type
2927
+ self.assertRaises(TypeError, C, '11', [0, 1], s)
2928
+ # count not sequence
2929
+ self.assertRaises(TypeError, C, a, {0, 1}, s)
2930
+ # count element not an int
2931
+ self.assertRaises(TypeError, C, a, [0, 1.0], s)
2932
+ # count element overflow
2933
+ self.assertRaises(OverflowError, C, a, [0, 1 << 65], s)
2934
+ # symbol not sequence
2935
+ self.assertRaises(TypeError, C, a, [0, 1], 43)
2936
+
2937
+ symbol = ['a', 'b', 'c', 'd']
2938
+ # sum(count) != len(symbol)
2939
+ self.assertRaisesMessage(ValueError,
2940
+ "sum(count) = 3, but len(symbol) = 4",
2941
+ C, a, [0, 1, 2], symbol)
2942
+ # count list too long
2943
+ self.assertRaisesMessage(ValueError,
2944
+ "len(count) cannot be larger than 32",
2945
+ C, a, 33 * [0], symbol)
2946
+
2947
+ def test_canonical_decode_count_range(self):
2948
+ a = bitarray()
2949
+ for i in range(1, 32):
2950
+ count = 32 * [0]
2951
+ # negative count
2952
+ count[i] = -1
2953
+ self.assertRaisesMessage(ValueError,
2954
+ "count[%d] not in [0..%d], got -1" % (i, 1 << i),
2955
+ canonical_decode, a, count, [])
2956
+
2957
+ maxbits = 1 << i
2958
+ count[i] = maxbits
2959
+ if i == 31 and PTRSIZE == 4:
2960
+ self.assertRaises(OverflowError,
2961
+ canonical_decode, a, count, [])
2962
+ continue
2963
+ self.assertRaisesMessage(ValueError,
2964
+ "sum(count) = %d, but len(symbol) = 0" % maxbits,
2965
+ canonical_decode, a, count, [])
2966
+
2967
+ count[i] = maxbits + 1
2968
+ self.assertRaisesMessage(ValueError,
2969
+ "count[%d] not in [0..%d], got %d" % (i, maxbits, count[i]),
2970
+ canonical_decode, a, count, [])
2971
+
2972
+ iter = canonical_decode(a, 32 * [0], [])
2973
+ self.assertEqual(list(iter), [])
2974
+
2975
+ def test_canonical_decode_simple(self):
2976
+ C = canonical_decode
2977
+ # symbols can be anything, they do not even have to be hashable here
2978
+ cnt = [0, 0, 4]
2979
+ s = ['A', 42, [1.2-3.7j, 4j], {'B': 6}]
2980
+ a = bitarray('00 01 10 11')
2981
+ # count can be a list
2982
+ self.assertEqual(list(C(a, cnt, s)), s)
2983
+ # count can also be a tuple (any sequence object in fact)
2984
+ self.assertEqual(list(C(a, (0, 0, 4), s)), s)
2985
+ self.assertEqual(list(C(7 * a, cnt, s)), 7 * s)
2986
+ # the count list may have extra 0's at the end (but not too many)
2987
+ count = [0, 0, 4, 0, 0, 0, 0, 0]
2988
+ self.assertEqual(list(C(a, count, s)), s)
2989
+ # the element count[0] is unused
2990
+ self.assertEqual(list(C(a, [-47, 0, 4], s)), s)
2991
+ # in fact it can be anything, as it is entirely ignored
2992
+ self.assertEqual(list(C(a, [None, 0, 4], s)), s)
2993
+
2994
+ # the symbol argument can be any sequence object
2995
+ s = [65, 66, 67, 98]
2996
+ self.assertEqual(list(C(a, cnt, s)), s)
2997
+ self.assertEqual(list(C(a, cnt, bytearray(s))), s)
2998
+ self.assertEqual(list(C(a, cnt, tuple(s))), s)
2999
+ self.assertEqual(list(C(a, cnt, bytes(s))), s)
3000
+ # Implementation Note:
3001
+ # The symbol can even be an iterable. This was done because we
3002
+ # want to use PySequence_Fast in order to convert sequence
3003
+ # objects (like bytes and bytearray) to a list. This is faster
3004
+ # as all objects are now elements in an array of pointers (as
3005
+ # opposed to having the object's __getitem__ method called on
3006
+ # every iteration).
3007
+ self.assertEqual(list(C(a, cnt, iter(s))), s)
3008
+
3009
+ def test_canonical_decode_empty(self):
3010
+ a = bitarray()
3011
+ # count and symbol are empty, ok because sum([]) == len([])
3012
+ self.assertEqual(list(canonical_decode(a, [], [])), [])
3013
+ a.append(0)
3014
+ self.assertRaisesMessage(ValueError, "reached end of bitarray",
3015
+ list, canonical_decode(a, [], []))
3016
+ a = bitarray(31 * '0')
3017
+ self.assertRaisesMessage(ValueError, "ran out of codes",
3018
+ list, canonical_decode(a, [], []))
3019
+
3020
+ def test_canonical_decode_one_symbol(self):
3021
+ symbols = ['A']
3022
+ count = [0, 1]
3023
+ a = bitarray('000')
3024
+ self.assertEqual(list(canonical_decode(a, count, symbols)),
3025
+ 3 * symbols)
3026
+ a.append(1)
3027
+ a.extend(bitarray(10 * '0'))
3028
+ iterator = canonical_decode(a, count, symbols)
3029
+ self.assertRaisesMessage(ValueError, "reached end of bitarray",
3030
+ list, iterator)
3031
+
3032
+ a.extend(bitarray(20 * '0'))
3033
+ iterator = canonical_decode(a, count, symbols)
3034
+ self.assertRaisesMessage(ValueError, "ran out of codes",
3035
+ list, iterator)
3036
+
3037
+ def test_canonical_decode_large(self):
3038
+ with open(__file__, 'rb') as f:
3039
+ msg = bytearray(f.read())
3040
+ self.assertTrue(len(msg) > 10_000)
3041
+ codedict, count, symbol = canonical_huffman(Counter(msg))
3042
+ a = bitarray()
3043
+ a.encode(codedict, msg)
3044
+ self.assertEqual(bytearray(canonical_decode(a, count, symbol)), msg)
3045
+ self.check_canonical_code(codedict, count, symbol)
3046
+
3047
+ def test_canonical_decode_symbol_change(self):
3048
+ msg = bytearray(b"Hello World!")
3049
+ codedict, count, symbol = canonical_huffman(Counter(msg))
3050
+ self.check_canonical_code(codedict, count, symbol)
3051
+ a = bitarray()
3052
+ a.encode(codedict, 10 * msg)
3053
+
3054
+ it = canonical_decode(a, count, symbol)
3055
+ def decode_one_msg():
3056
+ return bytearray(next(it) for _ in range(len(msg)))
3057
+
3058
+ self.assertEqual(decode_one_msg(), msg)
3059
+ symbol[symbol.index(ord("l"))] = ord("k")
3060
+ self.assertEqual(decode_one_msg(), bytearray(b"Hekko Workd!"))
3061
+ del symbol[:]
3062
+ self.assertRaises(IndexError, decode_one_msg)
3063
+
3064
+ def ensure_sorted(self, chc, symbol):
3065
+ # ensure codes are sorted
3066
+ for i in range(len(symbol) - 1):
3067
+ a = chc[symbol[i]]
3068
+ b = chc[symbol[i + 1]]
3069
+ self.assertTrue(ba2int(a) < ba2int(b))
3070
+
3071
+ def ensure_consecutive(self, chc, count, symbol):
3072
+ start = 0
3073
+ for nbits, cnt in enumerate(count):
3074
+ for i in range(start, start + cnt - 1):
3075
+ # ensure two consecutive codes (with same bit length) have
3076
+ # consecutive integer values
3077
+ a = chc[symbol[i]]
3078
+ b = chc[symbol[i + 1]]
3079
+ self.assertTrue(len(a) == len(b) == nbits)
3080
+ self.assertEqual(ba2int(a) + 1, ba2int(b))
3081
+ start += cnt
3082
+
3083
+ def ensure_count(self, chc, count):
3084
+ # ensure count list corresponds to length counts from codedict
3085
+ maxbits = len(count) - 1
3086
+ self.assertEqual(maxbits, max(len(a) for a in chc.values()))
3087
+ my_count = (maxbits + 1) * [0]
3088
+ for a in chc.values():
3089
+ self.assertEqual(a.endian, 'big')
3090
+ my_count[len(a)] += 1
3091
+ self.assertEqual(my_count, count)
3092
+
3093
+ def ensure_complete(self, count):
3094
+ # ensure code is complete and not oversubscribed
3095
+ maxbits = len(count) - 1
3096
+ x = sum(count[i] << (maxbits - i) for i in range(1, maxbits + 1))
3097
+ self.assertEqual(x, 1 << maxbits)
3098
+
3099
+ def ensure_round_trip(self, chc, count, symbol):
3100
+ # create a short test message, encode and decode
3101
+ msg = choices(symbol, k=10)
3102
+ a = bitarray()
3103
+ a.encode(chc, msg)
3104
+ it = canonical_decode(a, count, symbol)
3105
+ # the iterator holds a reference to the bitarray and symbol list
3106
+ del a, count, symbol
3107
+ self.assertEqual(type(it).__name__, 'canonical_decodeiter')
3108
+ self.assertEqual(list(it), msg)
3109
+
3110
+ def check_canonical_code(self, chc, count, symbol):
3111
+ self.check_code(chc)
3112
+ self.assertTrue(len(chc) == len(symbol) == sum(count))
3113
+ self.assertEqual(count[0], 0) # no codes have length 0
3114
+ self.assertEqual(set(chc), set(symbol))
3115
+ # the code of the last symbol has all 1 bits
3116
+ last = chc[symbol[-1]]
3117
+ self.assertTrue(last.all())
3118
+ # the code of the first symbol starts with bit 0
3119
+ first = chc[symbol[0]]
3120
+ self.assertEqual(first[0], 0)
3121
+
3122
+ self.ensure_sorted(chc, symbol)
3123
+ self.ensure_consecutive(chc, count, symbol)
3124
+ self.ensure_count(chc, count)
3125
+ self.ensure_complete(count)
3126
+ self.ensure_round_trip(chc, count, symbol)
3127
+
3128
+ def test_small_range(self):
3129
+ for n in range(2, 10):
3130
+ freq = {i: randrange(10) for i in range(n)}
3131
+ self.check_canonical_code(*canonical_huffman(freq))
3132
+
3133
+ def test_simple_counter(self):
3134
+ plain = bytearray(b'the quick brown fox jumps over the lazy dog.')
3135
+ cnt = Counter(plain)
3136
+ self.check_canonical_code(*canonical_huffman(cnt))
3137
+
3138
+ def test_no_comp(self):
3139
+ freq = {None: 1, "A": 1} # None and "A" are not comparable
3140
+ self.check_canonical_code(*canonical_huffman(freq))
3141
+
3142
+ def test_balanced(self):
3143
+ n = 7
3144
+ freq = {}
3145
+ for i in range(1 << n):
3146
+ freq[i] = 1
3147
+ code, count, sym = canonical_huffman(freq)
3148
+ self.assertEqual(len(code), 1 << n)
3149
+ self.assertTrue(all(len(v) == n for v in code.values()))
3150
+ self.check_canonical_code(code, count, sym)
3151
+
3152
+ def test_unbalanced(self):
3153
+ n = 32
3154
+ freq = {}
3155
+ for i in range(n):
3156
+ freq[i] = 1 << i
3157
+ code, count, sym = canonical_huffman(freq)
3158
+ self.assertEqual(len(code), n)
3159
+ for i in range(n):
3160
+ self.assertEqual(len(code[i]), n - max(1, i))
3161
+ self.check_canonical_code(code, count, sym)
3162
+
3163
+ # ---------------------------------------------------------------------------
3164
+
3165
+ if __name__ == '__main__':
3166
+ unittest.main()