1
0

bitfield.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. # The contents of this file are subject to the BitTorrent Open Source License
  2. # Version 1.1 (the License). You may not copy or use this file, in either
  3. # source code or executable form, except in compliance with the License. You
  4. # may obtain a copy of the License at http://www.bittorrent.com/license/.
  5. #
  6. # Software distributed under the License is distributed on an AS IS basis,
  7. # WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  8. # for the specific language governing rights and limitations under the
  9. # License.
  10. # Written by Bram Cohen, Uoti Urpala, and John Hoffman
  11. from array import array
  12. #counts = [chr(sum([(i >> j) & 1 for j in xrange(8)])) for i in xrange(256)]
  13. counts = []
  14. for i in xrange(256):
  15. t = 0
  16. for j in xrange(8):
  17. t = t + ((i >> j) & 1)
  18. counts.append(chr(t))
  19. counts = ''.join(counts)
  20. class Bitfield:
  21. def __init__(self, length, bitstring=None):
  22. self.length = length
  23. rlen, extra = divmod(length, 8)
  24. if bitstring is None:
  25. self.numfalse = length
  26. if extra:
  27. self.bits = array('B', chr(0) * (rlen + 1))
  28. else:
  29. self.bits = array('B', chr(0) * rlen)
  30. else:
  31. if extra:
  32. if len(bitstring) != rlen + 1:
  33. raise ValueError("%s != %s" % (len(bitstring), rlen + 1))
  34. if (ord(bitstring[-1]) << extra) & 0xFF != 0:
  35. raise ValueError("%s != %s" %
  36. ((ord(bitstring[-1]) << extra) & 0xFF, 0))
  37. else:
  38. if len(bitstring) != rlen:
  39. raise ValueError("%s != %s" % (len(bitstring), rlen))
  40. self.numfalse = length - sum(array('B',
  41. bitstring.translate(counts)))
  42. if self.numfalse != 0:
  43. self.bits = array('B', bitstring)
  44. else:
  45. self.bits = None
  46. def __setitem__(self, index, val):
  47. assert val
  48. pos = index >> 3
  49. mask = 128 >> (index & 7)
  50. if self.bits[pos] & mask:
  51. return
  52. self.bits[pos] = self.bits[pos] | mask
  53. self.numfalse = self.numfalse - 1
  54. if self.numfalse == 0:
  55. self.bits = None
  56. def __getitem__(self, index):
  57. bits = self.bits
  58. if bits is None:
  59. return 1
  60. return bits[index >> 3] & 128 >> (index & 7)
  61. def __len__(self):
  62. return self.length
  63. def tostring(self):
  64. if self.bits is None:
  65. rlen, extra = divmod(self.length, 8)
  66. r = chr(0xFF) * rlen
  67. if extra:
  68. r = r + chr((0xFF << (8 - extra)) & 0xFF)
  69. return r
  70. else:
  71. return self.bits.tostring()
  72. def __getstate__(self):
  73. d = {}
  74. d['length'] = self.length
  75. d['s'] = self.tostring()
  76. return d
  77. def __setstate__(self, d):
  78. Bitfield.__init__(self, d['length'], d['s'])
  79. old_Bitfield = Bitfield
  80. try:
  81. import BTL.cBitfield
  82. Bitfield = BTL.cBitfield.Bitfield
  83. except ImportError:
  84. pass