Raise error for unexpected dtype for Bit constructor [skip ci]

Author: ankane

pgvector/bit.py

@@ -7,9 +7,11 @@ def __init__(self, value):
         if isinstance(value, str):
             self._value = self.from_text(value)._value
         else:
-            # TODO raise if dtype not bool or uint8
-            if isinstance(value, np.ndarray) and value.dtype == np.uint8:
-                value = np.unpackbits(value)
+            if isinstance(value, np.ndarray):
+                if value.dtype == np.uint8:
+                    value = np.unpackbits(value).astype(bool)
+                elif value.dtype != np.bool:
+                    raise ValueError('expected dtype to be bool or uint8')
             else:
                 value = np.asarray(value, dtype=bool)
 

tests/test_bit.py

@@ -17,6 +17,12 @@ def test_ndarray_uint8(self):
         arr = np.array([254, 7, 0], dtype=np.uint8)
         assert Bit(arr).to_text() == '111111100000011100000000'
 
+    def test_ndarray_uint16(self):
+        arr = np.array([254, 7, 0], dtype=np.uint16)
+        with pytest.raises(ValueError) as error:
+            Bit(arr)
+        assert str(error.value) == 'expected dtype to be bool or uint8'
+
     def test_ndarray_same_object(self):
         arr = np.array([True, False, True])
         assert Bit(arr).to_list() == [True, False, True]