add export

Author: barrust

probables/utilities.py

@@ -4,6 +4,7 @@
 import mmap
 import string
 from array import array
+from io import IOBase
 from pathlib import Path
 from struct import Struct
 from typing import Literal, Union
@@ -214,3 +215,15 @@ def from_bytes(cls, data: bytes) -> "Bitarray":
         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

@@ -171,9 +171,22 @@ def test_bitarray_to_bytes(self):
             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))