Added check for ndim for Bit [skip ci]

ankane · ankane · commit db122fe5dc53 · 2024-06-01T19:31:53.000-07:00
diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py
@@ -9,7 +9,12 @@ def __init__(self, value):
         elif isinstance(value, str):
             self._value = __class__.from_text(value)._value
         else:
-            self._value = np.asarray(value, dtype=bool)
+            value = np.asarray(value, dtype=bool)
+
+            if value.ndim != 1:
+                raise ValueError('expected ndim to be 1')
+
+            self._value = value
 
     def __str__(self):
         return self.to_text()
diff --git a/tests/test_bit.py b/tests/test_bit.py
@@ -18,5 +18,15 @@ def test_ndarray_same_object(self):
         assert Bit(arr).to_list() == [True, False, True]
         assert Bit(arr).to_numpy() is arr
 
+    def test_ndim_two(self):
+        with pytest.raises(ValueError) as error:
+            Bit([[True, False], [True, False]])
+        assert str(error.value) == 'expected ndim to be 1'
+
+    def test_ndim_zero(self):
+        with pytest.raises(ValueError) as error:
+            Bit(True)
+        assert str(error.value) == 'expected ndim to be 1'
+
     def test_repr(self):
         assert repr(Bit([True, False, True])) == 'Bit(101)'
