Added type hints [skip ci]

Author: ankane

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.5.0 (unreleased)
 
+- Added type hints
 - Dropped support for Python < 3.10
 - Dropped support for SQLAlchemy < 2
 

pgvector/asyncpg/register.py

@@ -1,7 +1,8 @@
+from asyncpg import Connection
 from .. import Vector, HalfVector, SparseVector
 
 
-async def register_vector(conn, schema='public'):
+async def register_vector(conn: Connection, schema: str = 'public') -> None:
     await conn.set_type_codec(
         'vector',
         schema=schema,

pgvector/bit.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
 import numpy as np
 from struct import pack, unpack_from
+from typing import Any
 from warnings import warn
 
 
 class Bit:
-    def __init__(self, value):
+    def __init__(self, value: Any) -> None:
         if isinstance(value, bytes):
             self._len = 8 * len(value)
             self._data = value
@@ -26,32 +28,32 @@ def __init__(self, value):
             self._len = len(value)
             self._data = np.packbits(value).tobytes()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'Bit({self.to_text()})'
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         if isinstance(other, self.__class__):
             return self._len == other._len and self._data == other._data
         return False
 
-    def to_list(self):
+    def to_list(self) -> list[bool]:
         return self.to_numpy().tolist()
 
-    def to_numpy(self):
+    def to_numpy(self) -> np.ndarray:
         return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool)
 
-    def to_text(self):
+    def to_text(self) -> str:
         return ''.join(format(v, '08b') for v in self._data)[:self._len]
 
-    def to_binary(self):
+    def to_binary(self) -> bytes:
         return pack('>i', self._len) + self._data
 
     @classmethod
-    def from_text(cls, value):
+    def from_text(cls, value: str) -> Bit:
         return cls(str(value))
 
     @classmethod
-    def from_binary(cls, value):
+    def from_binary(cls, value: bytes) -> Bit:
         if not isinstance(value, bytes):
             raise ValueError('expected bytes')
 
@@ -61,15 +63,15 @@ def from_binary(cls, value):
         return bit
 
     @classmethod
-    def _to_db(cls, value):
-        if not isinstance(value, cls):
+    def _to_db(cls, value: Bit) -> str:
+        if not isinstance(value, Bit):
             raise ValueError('expected bit')
 
         return value.to_text()
 
     @classmethod
-    def _to_db_binary(cls, value):
-        if not isinstance(value, cls):
+    def _to_db_binary(cls, value: Bit) -> bytes:
+        if not isinstance(value, Bit):
             raise ValueError('expected bit')
 
         return value.to_binary()

pgvector/halfvec.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
 import numpy as np
 from struct import pack, unpack_from
+from typing import Any
 
 
 class HalfVector:
-    def __init__(self, value):
+    def __init__(self, value: Any) -> None:
         # asarray still copies if same dtype
         if not isinstance(value, np.ndarray) or value.dtype != '>f2':
             value = np.asarray(value, dtype='>f2')
@@ -13,40 +15,40 @@ def __init__(self, value):
 
         self._value = value
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f'HalfVector({self.to_list()})'
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         if isinstance(other, self.__class__):
             return np.array_equal(self.to_numpy(), other.to_numpy())
         return False
 
-    def dimensions(self):
+    def dimensions(self) -> int:
         return len(self._value)
 
-    def to_list(self):
+    def to_list(self) -> list[float]:
         return self._value.tolist()
 
-    def to_numpy(self):
+    def to_numpy(self) -> np.ndarray:
         return self._value
 
-    def to_text(self):
+    def to_text(self) -> str:
         return '[' + ','.join([str(float(v)) for v in self._value]) + ']'
 
-    def to_binary(self):
+    def to_binary(self) -> bytes:
         return pack('>HH', self.dimensions(), 0) + self._value.tobytes()
 
     @classmethod
-    def from_text(cls, value):
+    def from_text(cls, value: str) -> HalfVector:
         return cls([float(v) for v in value[1:-1].split(',')])
 
     @classmethod
-    def from_binary(cls, value):
+    def from_binary(cls, value: bytes) -> HalfVector:
         dim, unused = unpack_from('>HH', value)
         return cls(np.frombuffer(value, dtype='>f2', count=dim, offset=4))
 
     @classmethod
-    def _to_db(cls, value, dim=None):
+    def _to_db(cls, value: Any, dim: int | None = None) -> str | None:
         if value is None:
             return value
 
@@ -59,7 +61,7 @@ def _to_db(cls, value, dim=None):
         return value.to_text()
 
     @classmethod
-    def _to_db_binary(cls, value):
+    def _to_db_binary(cls, value: Any) -> bytes | None:
         if value is None:
             return value
 
@@ -69,15 +71,15 @@ def _to_db_binary(cls, value):
         return value.to_binary()
 
     @classmethod
-    def _from_db(cls, value):
-        if value is None or isinstance(value, cls):
+    def _from_db(cls, value: str | HalfVector | None) -> HalfVector | None:
+        if value is None or isinstance(value, HalfVector):
             return value
 
         return cls.from_text(value)
 
     @classmethod
-    def _from_db_binary(cls, value):
-        if value is None or isinstance(value, cls):
+    def _from_db_binary(cls, value: bytes | HalfVector | None) -> HalfVector | None:
+        if value is None or isinstance(value, HalfVector):
             return value
 
         return cls.from_binary(value)

pgvector/pg8000/register.py

@@ -1,8 +1,9 @@
 import numpy as np
+from pg8000.native import Connection
 from .. import Vector, HalfVector, SparseVector
 
 
-def register_vector(conn):
+def register_vector(conn: Connection) -> None:
     # use to_regtype to get first matching type in search path
     res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")
     type_info = dict(res)

pgvector/psycopg/bit.py

@@ -1,25 +1,30 @@
+from psycopg import BaseConnection
+from psycopg.types import TypeInfo
 from psycopg.adapt import Dumper
 from psycopg.pq import Format
+from typing import Any, TypeAlias
 from .. import Bit
 
+Buffer: TypeAlias = bytes | bytearray | memoryview
+
 
 class BitDumper(Dumper):
 
     format = Format.TEXT
 
-    def dump(self, obj):
+    def dump(self, obj: Bit) -> Buffer | None:
         return Bit._to_db(obj).encode('utf8')
 
 
 class BitBinaryDumper(BitDumper):
 
     format = Format.BINARY
 
-    def dump(self, obj):
+    def dump(self, obj: Bit) -> Buffer | None:
         return Bit._to_db_binary(obj)
 
 
-def register_bit_info(context, info):
+def register_bit_info(context: BaseConnection[Any], info: TypeInfo) -> None:
     info.register(context)
 
     # add oid to anonymous class for set_types

pgvector/psycopg/halfvec.py

@@ -1,29 +1,35 @@
+from psycopg import BaseConnection
 from psycopg.adapt import Loader, Dumper
 from psycopg.pq import Format
+from psycopg.types import TypeInfo
+from typing import Any, TypeAlias
 from .. import HalfVector
 
+Buffer: TypeAlias = bytes | bytearray | memoryview
+
 
 class HalfVectorDumper(Dumper):
 
     format = Format.TEXT
 
-    def dump(self, obj):
-        return HalfVector._to_db(obj).encode('utf8')
+    def dump(self, obj: HalfVector) -> Buffer | None:
+        value = HalfVector._to_db(obj)
+        return value if value is None else value.encode('utf8')
 
 
 class HalfVectorBinaryDumper(HalfVectorDumper):
 
     format = Format.BINARY
 
-    def dump(self, obj):
+    def dump(self, obj: HalfVector) -> Buffer | None:
         return HalfVector._to_db_binary(obj)
 
 
 class HalfVectorLoader(Loader):
 
     format = Format.TEXT
 
-    def load(self, data):
+    def load(self, data: Buffer) -> HalfVector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
         return HalfVector._from_db(data.decode('utf8'))
@@ -33,13 +39,13 @@ class HalfVectorBinaryLoader(HalfVectorLoader):
 
     format = Format.BINARY
 
-    def load(self, data):
-        if isinstance(data, memoryview):
+    def load(self, data: Buffer) -> HalfVector | None:
+        if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
         return HalfVector._from_db_binary(data)
 
 
-def register_halfvec_info(context, info):
+def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo) -> None:
     info.register(context)
 
     # add oid to anonymous class for set_types

pgvector/psycopg/register.py

@@ -1,15 +1,18 @@
+from psycopg import AsyncConnection, Connection
 from psycopg.types import TypeInfo
+from typing import Any
 from .bit import register_bit_info
 from .halfvec import register_halfvec_info
 from .sparsevec import register_sparsevec_info
 from .vector import register_vector_info
 
 
-def register_vector(context):
+def register_vector(context: Connection[Any]) -> None:
     info = TypeInfo.fetch(context, 'vector')
     register_vector_info(context, info)
 
     info = TypeInfo.fetch(context, 'bit')
+    assert info is not None
     register_bit_info(context, info)
 
     info = TypeInfo.fetch(context, 'halfvec')
@@ -21,11 +24,12 @@ def register_vector(context):
         register_sparsevec_info(context, info)
 
 
-async def register_vector_async(context):
+async def register_vector_async(context: AsyncConnection[Any]) -> None:
     info = await TypeInfo.fetch(context, 'vector')
     register_vector_info(context, info)
 
     info = await TypeInfo.fetch(context, 'bit')
+    assert info is not None
     register_bit_info(context, info)
 
     info = await TypeInfo.fetch(context, 'halfvec')

pgvector/psycopg/sparsevec.py

@@ -1,29 +1,35 @@
+from psycopg import BaseConnection
 from psycopg.adapt import Loader, Dumper
 from psycopg.pq import Format
+from psycopg.types import TypeInfo
+from typing import Any, TypeAlias
 from .. import SparseVector
 
+Buffer: TypeAlias = bytes | bytearray | memoryview
+
 
 class SparseVectorDumper(Dumper):
 
     format = Format.TEXT
 
-    def dump(self, obj):
-        return SparseVector._to_db(obj).encode('utf8')
+    def dump(self, obj: SparseVector) -> Buffer | None:
+        value = SparseVector._to_db(obj)
+        return value if value is None else value.encode('utf8')
 
 
 class SparseVectorBinaryDumper(SparseVectorDumper):
 
     format = Format.BINARY
 
-    def dump(self, obj):
+    def dump(self, obj: SparseVector) -> Buffer | None:
         return SparseVector._to_db_binary(obj)
 
 
 class SparseVectorLoader(Loader):
 
     format = Format.TEXT
 
-    def load(self, data):
+    def load(self, data: Buffer) -> SparseVector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
         return SparseVector._from_db(data.decode('utf8'))
@@ -33,13 +39,13 @@ class SparseVectorBinaryLoader(SparseVectorLoader):
 
     format = Format.BINARY
 
-    def load(self, data):
-        if isinstance(data, memoryview):
+    def load(self, data: Buffer) -> SparseVector | None:
+        if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
         return SparseVector._from_db_binary(data)
 
 
-def register_sparsevec_info(context, info):
+def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo) -> None:
     info.register(context)
 
     # add oid to anonymous class for set_types