Added test for Bit constructor with uint8

pgvector · Jul 17, 2024 · 4498caa · 4498caa
1 parent e77ee13
commit 4498caa
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py
@@ -7,7 +7,12 @@ def __init__(self, value):
         if isinstance(value, str):
             self._value = self.from_text(value)._value
         else:
-            # TODO use np.unpackbits for uint8 array
+            # TODO change in 0.4.0
+            # if isinstance(value, np.ndarray) and value.dtype == np.uint8:
+            #     value = np.unpackbits(value)
+            # else:
+            #     value = np.asarray(value, dtype=bool)
+
             value = np.asarray(value, dtype=bool)
 
             if value.ndim != 1:

diff --git a/tests/test_bit.py b/tests/test_bit.py
@@ -13,6 +13,12 @@ def test_tuple(self):
     def test_str(self):
         assert Bit('101').to_list() == [True, False, True]
 
+    def test_ndarray_uint8(self):
+        arr = np.array([254, 7, 0], dtype=np.uint8)
+        # TODO change in 0.4.0
+        # assert Bit(arr).to_text() == '111111100000011100000000'
+        assert Bit(arr).to_text() == '110'
+
     def test_ndarray_same_object(self):
         arr = np.array([True, False, True])
         assert Bit(arr).to_list() == [True, False, True]