Added typing to attributes

Author: ankane

pgvector/bit.py

@@ -8,6 +8,9 @@
 
 
 class Bit:
+    _length: int
+    _data: bytes
+
     def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]]) -> None:
         if isinstance(value, bytes):
             self._length = 8 * len(value)

pgvector/halfvec.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import array
+from array import array
 import struct
 import sys
 from typing import TYPE_CHECKING
@@ -10,11 +10,13 @@
 
 
 class HalfVector:
+    _value: array[int]
+
     def __init__(self, value: list[float] | ndarray) -> None:
         if isinstance(value, list):
             dim = len(value)
             try:
-                self._value = array.array('H', struct.pack(f'{dim}e', *value))
+                self._value = array('H', struct.pack(f'{dim}e', *value))
             except struct.error:
                 raise ValueError('expected list[float]')
         elif is_ndarray(value):
@@ -23,7 +25,7 @@ def __init__(self, value: list[float] | ndarray) -> None:
             if value.ndim != 1:
                 raise ValueError('expected ndim to be 1')
 
-            arr = array.array('H')
+            arr = array('H')
             arr.frombytes(value.astype(np.float16, order='C', copy=False).data.cast('B'))
             self._value = arr
         else:
@@ -55,7 +57,7 @@ def to_binary(self) -> bytes:
         if sys.byteorder == 'big':
             value = self._value
         else:
-            value = array.array('H', self._value)
+            value = array('H', self._value)
             value.byteswap()
         return struct.pack('>HH', len(value), 0) + memoryview(value)
 
@@ -74,7 +76,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> HalfVector:
         if unused != 0:
             raise ValueError('expected unused to be 0')
 
-        arr = array.array('H')
+        arr = array('H')
         arr.frombytes(data)
         if sys.byteorder != 'big':
             arr.byteswap()

pgvector/sparsevec.py

@@ -16,6 +16,10 @@ class Sentinel:
 
 
 class SparseVector:
+    _dim: int
+    _indices: list[int]
+    _values: list[float]
+
     @overload
     def __init__(self, value: dict[int, float], dimensions: int, /) -> None:
         ...

pgvector/vector.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import array
+from array import array
 import struct
 import sys
 from typing import TYPE_CHECKING
@@ -10,10 +10,12 @@
 
 
 class Vector:
+    _value: array[float]
+
     def __init__(self, value: list[float] | ndarray) -> None:
         if isinstance(value, list):
             try:
-                self._value = array.array('f', value)
+                self._value = array('f', value)
             except TypeError:
                 raise ValueError('expected list[float]')
         elif is_ndarray(value):
@@ -22,7 +24,7 @@ def __init__(self, value: list[float] | ndarray) -> None:
             if value.ndim != 1:
                 raise ValueError('expected ndim to be 1')
 
-            arr = array.array('f')
+            arr = array('f')
             arr.frombytes(value.astype(np.float32, order='C', copy=False).data.cast('B'))
             self._value = arr
         else:
@@ -53,7 +55,7 @@ def to_binary(self) -> bytes:
         if sys.byteorder == 'big':
             value = self._value
         else:
-            value = array.array('f', self._value)
+            value = array('f', self._value)
             value.byteswap()
         return struct.pack('>HH', len(value), 0) + memoryview(value)
 
@@ -72,7 +74,7 @@ def from_binary(cls, value: bytes | bytearray | memoryview) -> Vector:
         if unused != 0:
             raise ValueError('expected unused to be 0')
 
-        arr = array.array('f')
+        arr = array('f')
         arr.frombytes(data)
         if sys.byteorder != 'big':
             arr.byteswap()