Fix bitarray slice result

mlouielu · Jul 3, 2017 · d61b748 · d61b748
1 parent 806e40f
commit d61b748
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/nand2vm/bitarray.py b/nand2vm/bitarray.py
@@ -29,6 +29,8 @@ def __init__(self, source, endian=True):
             self.data.reverse()
 
     def __getitem__(self, key):
+        if isinstance(key, slice):
+            return BitArray(self.data[key], endian=False)
         return self.data[key]
 
     def __setitem__(self, key, value):

diff --git a/test/test_bitarray.py b/test/test_bitarray.py
@@ -42,3 +42,13 @@ def test_init_with_int(self):
 
     def test_len(self):
         self.assertEqual(len(BitArray('110')), 3)
+
+    def test_slice(self):
+        source = [True, True, False, True]
+        b = BitArray(source)
+
+        self.assertIsInstance(b[:2], BitArray)
+        self.assertEqual(b[:2], BitArray([False, True]))
+
+        self.assertIsInstance(b[2:], BitArray)
+        self.assertEqual(b[2:], BitArray([True, True]))