78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
# 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()
|