bitarray export

probables/utilities.py

@@ -4,8 +4,10 @@
 import mmap
 import string
 from array import array
+from io import IOBase
 from pathlib import Path
-from typing import Union
+from struct import Struct
+from typing import Literal, Union
 
 
 def is_hex_string(hex_string: Union[str, None]) -> bool:
@@ -74,7 +76,7 @@ def close(self) -> None:
         """Close the MMap class includeing cleaning up open files, etc"""
         self.__exit__()
 
-    def seek(self, pos: int, whence: int) -> None:
+    def seek(self, pos: int, whence: Literal[0, 1, 2]) -> None:
         """Implement a method to seek on top of the MMap class"""
         self.__m.seek(pos, whence)
 
@@ -190,3 +192,38 @@ def num_bits_set(self) -> int:
         Returns:
             int: Number of bits set"""
         return sum(self.check_bit(x) for x in range(self._size))
+
+    _BITARRAY_FOOTER = Struct("Q")
+    _BITARRAY_BIN = Struct("B")
+
+    def to_bytes(self) -> bytes:
+        """Convert the bitarray to bytes
+
+        Returns:
+            bytes: Bitarray representation as bytes"""
+        footer = self._BITARRAY_FOOTER.pack(self._size)
+        return self._bitarray.tobytes() + footer
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "Bitarray":
+        """Convert bytes to a bitarray
+
+        Args:
+            data (bytes): Bytes to convert"""
+        size = cls._BITARRAY_FOOTER.unpack(data[-8:])[0]
+        bitarray = array("B", data[:-8])
+        ba = Bitarray(size)
+        ba._bitarray = bitarray
+        return ba
+
+    def export(self, file: Union[Path, str, IOBase, mmap.mmap]) -> None:
+        """Export the bitarray to a file
+
+        Args:
+            filename (str): Filename to export to"""
+        if not isinstance(file, (IOBase, mmap.mmap)):
+            file = resolve_path(file)
+            with open(file, "wb") as filepointer:
+                self.export(filepointer)
+        else:
+            file.write(self.to_bytes())

tests/test_utilities.py

@@ -159,9 +159,34 @@ def test_bitarray(self):
             "0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
         )
 
+    def test_bitarray_to_bytes(self):
+        ba = Bitarray(100)
+        for i in range(34):
+            ba.set_bit(i * 3)
+        b = ba.to_bytes()
+        nba = Bitarray.from_bytes(b)
+        self.assertEqual(nba.as_string(), ba.as_string())
+        for i in range(100):
+            self.assertEqual(ba[i], nba[i])
+            print(i)
+            self.assertEqual(nba[i], 0 if i % 3 != 0 else 1)
+
+    def test_bitarray_export(self):
+        with NamedTemporaryFile() as tmp:
+            ba = Bitarray(100)
+            for i in range(34):
+                ba.set_bit(i * 3)
+            ba.export(tmp.name)
+            nba = Bitarray.from_bytes(tmp.read())
+            self.assertEqual(nba.as_string(), ba.as_string())
+            for i in range(100):
+                self.assertEqual(ba[i], nba[i])
+                print(i)
+                self.assertEqual(nba[i], 0 if i % 3 != 0 else 1)
+
     def test_bitarray_invalid_idx(self):
         """use an invalid type in a jaccard index"""
-        self.assertRaises(TypeError, lambda: Bitarray("100"))
+        self.assertRaises(TypeError, lambda: Bitarray("100"))  # type: ignore
         self.assertRaises(ValueError, lambda: Bitarray(-100))
         ba = Bitarray(10)
         self.assertRaises(IndexError, lambda: ba.set_bit(12))