1 files changed, 77 insertions, 0 deletions
diff --git a/BitTorrent/bitfield.py b/BitTorrent/bitfield.py
new file mode 100644
index 0000000..f9fb511
--- /dev/null
+++ b/BitTorrent/bitfield.py
@@ -0,0 +1,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()