diff --git a/pgvector/bit.py b/pgvector/bit.py index 098ebb2..a9370e4 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -88,7 +88,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit: length, = unpack_from('>i', value) data = memoryview(value)[4:].tobytes() - if len(data) != (length + 7) // 8: + if length < 0 or len(data) != (length + 7) // 8: raise ValueError('invalid length') bit = cls.__new__(cls) diff --git a/tests/test_bit.py b/tests/test_bit.py index 6487348..ad1bd36 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -6,6 +6,16 @@ class TestBit: + @pytest.mark.parametrize('length', [-1, -2, -7, -8]) + def test_from_binary_negative_length(self, length: int) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(pack('>i', length)) + + @pytest.mark.parametrize('length', [0, 1, 7, 8, 9]) + def test_from_binary_length_boundaries(self, length: int) -> None: + bit = Bit('1' * length) + assert Bit.from_binary(bit.to_binary()) == bit + def test_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True]