Improved type hints [skip ci]

Author: ankane

pgvector/sparsevec.py

@@ -42,7 +42,7 @@ def indices(self) -> list[int]:
     def values(self) -> list[float]:
         return self._values
 
-    def to_coo(self):
+    def to_coo(self) -> Any:
         from scipy.sparse import coo_array
 
         coords = ([0] * len(self._indices), self._indices)

pgvector/sqlalchemy/bit.py

@@ -20,7 +20,7 @@ def bind_processor(self, dialect: Dialect) -> Any:
         if dialect.__class__.__name__ == 'PGDialect_asyncpg':
             import asyncpg
 
-            def process(value):
+            def process(value: Any) -> Any:
                 if isinstance(value, str):
                     return asyncpg.BitString(value)
                 return value

pgvector/sqlalchemy/halfvec.py

@@ -19,19 +19,19 @@ def get_col_spec(self, **kw: Any) -> str:
         return 'HALFVEC(%d)' % self.dim
 
     def bind_processor(self, dialect: Dialect) -> Any:
-        def process(value):
+        def process(value: Any) -> str | None:
             return HalfVector._to_db(value, self.dim)
         return process
 
     def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
-        def process(value):
+        def process(value: Any) -> Any:
             return string_literal_processor(HalfVector._to_db(value, self.dim))  # type: ignore
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value):
+        def process(value: Any) -> HalfVector | None:
             return HalfVector._from_db(value)
         return process
 

pgvector/sqlalchemy/sparsevec.py

@@ -19,19 +19,19 @@ def get_col_spec(self, **kw: Any) -> str:
         return 'SPARSEVEC(%d)' % self.dim
 
     def bind_processor(self, dialect: Dialect) -> Any:
-        def process(value):
+        def process(value: Any) -> str | None:
             return SparseVector._to_db(value, self.dim)
         return process
 
     def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
-        def process(value):
+        def process(value: Any) -> Any:
             return string_literal_processor(SparseVector._to_db(value, self.dim))  # type: ignore
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value):
+        def process(value: Any) -> SparseVector | None:
             return SparseVector._from_db(value)
         return process
 

pgvector/sqlalchemy/vector.py

@@ -1,3 +1,4 @@
+import numpy as np
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float, String
 from sqlalchemy import Dialect, Operators
@@ -19,19 +20,19 @@ def get_col_spec(self, **kw: Any) -> str:
         return 'VECTOR(%d)' % self.dim
 
     def bind_processor(self, dialect: Dialect) -> Any:
-        def process(value):
+        def process(value: Any) -> str | None:
             return Vector._to_db(value, self.dim)
         return process
 
     def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
-        def process(value):
+        def process(value: Any) -> Any:
             return string_literal_processor(Vector._to_db(value, self.dim))  # type: ignore
         return process
 
     def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
-        def process(value):
+        def process(value: Any) -> np.ndarray | None:
             return Vector._from_db(value)
         return process