Merge branch 'master' into numpy-dep

Author: ankane

pgvector/bit.py

@@ -2,23 +2,39 @@
 from struct import pack, unpack_from
 from warnings import warn
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class Bit:
     def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None:
         if isinstance(value, bytes):
             length = 8 * len(value)
             data = value
         else:
-            import numpy as np
+            if isinstance(value, list):
+                def bit_value(v: bool) -> str:
+                    if v is True:
+                        return '1'
+                    if v is False:
+                        return '0'
+                    raise ValueError('expected list[bool]')
+
+                value = ''.join([bit_value(v) for v in value])
 
             if isinstance(value, str):
-                value = [v != '0' for v in value]
-            else:
-                value = np.asarray(value)
+                length = len(value)
 
-                # for mypy
-                assert isinstance(value, np.ndarray)
+                if length % 8 != 0:
+                    value += '0' * (8 - (length % 8))
 
+                try:
+                    data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
+                except ValueError:
+                    raise ValueError('expected bit string')
+            elif np is not None and isinstance(value, np.ndarray):
                 if value.dtype != np.bool:
                     # skip warning for result of np.unpackbits
                     if value.dtype != np.uint8 or np.any(value > 1):
@@ -31,8 +47,10 @@ def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.d
                 if value.ndim != 1:
                     raise ValueError('expected ndim to be 1')
 
-            length = len(value)
-            data = np.packbits(value).tobytes()
+                length = len(value)
+                data = np.packbits(value).tobytes()
+            else:
+                raise ValueError('expected bytes, str, list, or ndarray')
 
         self._value = pack('>i', length) + data
 

tests/test_bit.py

@@ -8,24 +8,27 @@
 
 
 class TestBit:
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_list(self):
         assert Bit([True, False, True]).to_list() == [True, False, True]
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_list_none(self):
-        with pytest.warns(UserWarning, match='expected elements to be boolean'):
-            assert Bit([True, None, True]).to_text() == '101'  # ty: ignore[invalid-argument-type]
+        with pytest.raises(ValueError) as error:
+            Bit([True, None, True])  # ty: ignore[invalid-argument-type]
+        assert str(error.value) == 'expected list[bool]'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_list_int(self):
-        with pytest.warns(UserWarning, match='expected elements to be boolean'):
-            assert Bit([254, 7, 0]).to_text() == '110'  # ty: ignore[invalid-argument-type]
+        with pytest.raises(ValueError) as error:
+            Bit([254, 7, 0])  # ty: ignore[invalid-argument-type]
+        assert str(error.value) == 'expected list[bool]'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_str(self):
         assert Bit('101').to_list() == [True, False, True]
 
+    def test_str_two(self):
+        with pytest.raises(ValueError) as error:
+            Bit('201')
+        assert str(error.value) == 'expected bit string'
+
     def test_bytes(self):
         assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000'
         assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000'
@@ -53,24 +56,20 @@ def test_ndarray_uint16(self):
         with pytest.warns(UserWarning, match='expected elements to be boolean'):
             assert Bit(arr).to_text() == '110'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
             Bit([[True, False], [True, False]])  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected list[bool]'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndim_zero(self):
         with pytest.raises(ValueError) as error:
             Bit(True)  # ty: ignore[invalid-argument-type]
-        assert str(error.value) == 'expected ndim to be 1'
+        assert str(error.value) == 'expected bytes, str, list, or ndarray'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_repr(self):
         assert repr(Bit([True, False, True])) == 'Bit(101)'
         assert str(Bit([True, False, True])) == 'Bit(101)'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_equality(self):
         assert Bit([True, False, True]) == Bit([True, False, True])
         assert Bit([True, False, True]) != Bit([True, False, False])

tests/test_psycopg.py

@@ -90,21 +90,18 @@ def test_halfvec_text_format(self):
         res = next(conn.execute('SELECT %t::halfvec', (embedding,)))[0]
         assert res == HalfVector([1.5, 2, 3])
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_bit(self):
         embedding = Bit([True, False, True])
         conn.execute('INSERT INTO psycopg_items (binary_embedding) VALUES (%s)', (embedding,))
 
         res = next(conn.execute('SELECT binary_embedding FROM psycopg_items ORDER BY id'))[0]
         assert res == '101'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_bit_binary_format(self):
         embedding = Bit([False, True, False, True, False, False, False, False, True])
         res = next(conn.execute('SELECT %b::bit(9)', (embedding,), binary=True))[0]
         assert repr(Bit.from_binary(res)) == 'Bit(010100001)'
 
-    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_bit_text_format(self):
         embedding = Bit([False, True, False, True, False, False, False, False, True])
         res = next(conn.execute('SELECT %t::bit(9)', (embedding,)))[0]