1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
import unittest
from random import getrandbits, randint, randrange
from time import perf_counter
from bitarray import bitarray
from bitarray.util import urandom
def nxir(x, start, step):
assert x >= start and step > 0
# in Python we can use a simpler expression than in C
return (start - x) % step
def set_range_opt(self, start, stop, step, value):
ca = (start + 7) // 8
cb = stop // 8
m = (cb - ca) * 8
assert m >= 0
assert 0 <= 8 * ca - start < 8
assert 0 <= stop - 8 * cb < 8
mask = bitarray(step, self.endian)
mask.setall(not value)
mask[nxir(8 * ca, start, step)] = value
mask *= (m - 1) // step + 1
del mask[m:] # in the C version we wouldn't bother
assert len(mask) % 8 == 0
self[start : 8 * ca : step] = value
if value:
self[8 * ca : 8 * cb] |= mask
else:
self[8 * ca : 8 * cb] &= mask
self[8 * cb + nxir(8 * cb, start, step) : stop : step] = value
class Tests(unittest.TestCase):
def test_nxir(self):
for _ in range(1000):
start = randrange(100)
step = randrange(1, 20)
x = randrange(start, start + 100)
nx = nxir(x, start, step)
self.assertTrue(0 <= nx < step)
self.assertEqual((x + nx) % step, start % step)
def test_setslice_bool_step(self):
# this test exercises set_range() when stop is much larger than start
for _ in range(5000):
n = randrange(3000, 4000)
a = urandom(n)
aa = a.tolist()
s = slice(randrange(1000), randrange(1000, n), randint(1, 100))
self.assertTrue(s.stop - s.start >= 0)
slicelength = len(range(n)[s])
self.assertTrue(slicelength > 0)
v = getrandbits(1)
set_range_opt(a, s.start, s.stop, s.step, v)
aa[s] = slicelength * [v]
self.assertEqual(a.tolist(), aa)
def speed_cmp():
n = 1_000_000
print("n=%d\ntimes in micro seconds\n" % n)
print('%8s %12s %12s' % ("step", "this-code", "native"))
for step in range(1, 20):
a = bitarray(n)
b = bitarray(n)
t0 = perf_counter()
set_range_opt(a, 0, n, step, 1)
t1 = perf_counter()
b[::step] = 1
t2 = perf_counter()
print('%8d %12.3f %12.3f' % (step, 1E6 * (t1 - t0), 1E6 * (t2 - t1)))
assert a == b
if __name__ == '__main__':
speed_cmp()
unittest.main()
|
