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
|
# The contents of this file are subject to the BitTorrent Open Source License
# Version 1.1 (the License). You may not copy or use this file, in either
# source code or executable form, except in compliance with the License. You
# may obtain a copy of the License at http://www.bittorrent.com/license/.
#
# Software distributed under the License is distributed on an AS IS basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
# Written by Bram Cohen, Uoti Urpala, and John Hoffman
from array import array
from BitTorrent.obsoletepythonsupport import *
counts = [chr(sum([(i >> j) & 1 for j in xrange(8)])) for i in xrange(256)]
counts = ''.join(counts)
class Bitfield:
def __init__(self, length, bitstring=None):
self.length = length
rlen, extra = divmod(length, 8)
if bitstring is None:
self.numfalse = length
if extra:
self.bits = array('B', chr(0) * (rlen + 1))
else:
self.bits = array('B', chr(0) * rlen)
else:
if extra:
if len(bitstring) != rlen + 1:
raise ValueError
if (ord(bitstring[-1]) << extra) & 0xFF != 0:
raise ValueError
else:
if len(bitstring) != rlen:
raise ValueError
c = counts
self.numfalse = length - sum(array('B',
bitstring.translate(counts)))
if self.numfalse != 0:
self.bits = array('B', bitstring)
else:
self.bits = None
def __setitem__(self, index, val):
assert val
pos = index >> 3
mask = 128 >> (index & 7)
if self.bits[pos] & mask:
return
self.bits[pos] |= mask
self.numfalse -= 1
if self.numfalse == 0:
self.bits = None
def __getitem__(self, index):
bits = self.bits
if bits is None:
return 1
return bits[index >> 3] & 128 >> (index & 7)
def __len__(self):
return self.length
def tostring(self):
if self.bits is None:
rlen, extra = divmod(self.length, 8)
r = chr(0xFF) * rlen
if extra:
r += chr((0xFF << (8 - extra)) & 0xFF)
return r
else:
return self.bits.tostring()
|