Added more type hints [skip ci]

Author: ankane

pgvector/sqlalchemy/bit.py

@@ -1,20 +1,22 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float
+from sqlalchemy import Dialect
+from typing import Any
 
 
 class BIT(UserDefinedType):
     cache_ok = True
 
-    def __init__(self, length=None):
+    def __init__(self, length: int | None = None) -> None:
         super(UserDefinedType, self).__init__()
         self.length = length
 
-    def get_col_spec(self, **kw):
+    def get_col_spec(self, **kw) -> str:
         if self.length is None:
             return 'BIT'
         return 'BIT(%d)' % self.length
 
-    def bind_processor(self, dialect):
+    def bind_processor(self, dialect: Dialect) -> Any:
         if dialect.__class__.__name__ == 'PGDialect_asyncpg':
             import asyncpg
 
@@ -27,10 +29,10 @@ def process(value):
             return super().bind_processor(dialect)
 
     class comparator_factory(UserDefinedType.Comparator):
-        def hamming_distance(self, other):
+        def hamming_distance(self, other: Any) -> Any:
             return self.op('<~>', return_type=Float)(other)
 
-        def jaccard_distance(self, other):
+        def jaccard_distance(self, other: Any) -> Any:
             return self.op('<%>', return_type=Float)(other)
 
 

pgvector/sqlalchemy/halfvec.py

@@ -1,49 +1,51 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy import Dialect
+from typing import Any
 from .. import HalfVector
 
 
 class HALFVEC(UserDefinedType):
     cache_ok = True
     _string = String()
 
-    def __init__(self, dim=None):
+    def __init__(self, dim: int | None = None) -> None:
         super(UserDefinedType, self).__init__()
         self.dim = dim
 
-    def get_col_spec(self, **kw):
+    def get_col_spec(self, **kw) -> str:
         if self.dim is None:
             return 'HALFVEC'
         return 'HALFVEC(%d)' % self.dim
 
-    def bind_processor(self, dialect):
+    def bind_processor(self, dialect: Dialect):
         def process(value):
             return HalfVector._to_db(value, self.dim)
         return process
 
-    def literal_processor(self, dialect):
+    def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
         def process(value):
             return string_literal_processor(HalfVector._to_db(value, self.dim))  # type: ignore
         return process
 
-    def result_processor(self, dialect, coltype):
+    def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
         def process(value):
             return HalfVector._from_db(value)
         return process
 
     class comparator_factory(UserDefinedType.Comparator):
-        def l2_distance(self, other):
+        def l2_distance(self, other: Any) -> Any:
             return self.op('<->', return_type=Float)(other)
 
-        def max_inner_product(self, other):
+        def max_inner_product(self, other: Any) -> Any:
             return self.op('<#>', return_type=Float)(other)
 
-        def cosine_distance(self, other):
+        def cosine_distance(self, other: Any) -> Any:
             return self.op('<=>', return_type=Float)(other)
 
-        def l1_distance(self, other):
+        def l1_distance(self, other: Any) -> Any:
             return self.op('<+>', return_type=Float)(other)
 
 

pgvector/sqlalchemy/sparsevec.py

@@ -1,49 +1,51 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy import Dialect
+from typing import Any
 from .. import SparseVector
 
 
 class SPARSEVEC(UserDefinedType):
     cache_ok = True
     _string = String()
 
-    def __init__(self, dim=None):
+    def __init__(self, dim: int | None = None) -> None:
         super(UserDefinedType, self).__init__()
         self.dim = dim
 
-    def get_col_spec(self, **kw):
+    def get_col_spec(self, **kw) -> str:
         if self.dim is None:
             return 'SPARSEVEC'
         return 'SPARSEVEC(%d)' % self.dim
 
-    def bind_processor(self, dialect):
+    def bind_processor(self, dialect: Dialect) -> Any:
         def process(value):
             return SparseVector._to_db(value, self.dim)
         return process
 
-    def literal_processor(self, dialect):
+    def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
         def process(value):
             return string_literal_processor(SparseVector._to_db(value, self.dim))  # type: ignore
         return process
 
-    def result_processor(self, dialect, coltype):
+    def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
         def process(value):
             return SparseVector._from_db(value)
         return process
 
     class comparator_factory(UserDefinedType.Comparator):
-        def l2_distance(self, other):
+        def l2_distance(self, other: Any) -> Any:
             return self.op('<->', return_type=Float)(other)
 
-        def max_inner_product(self, other):
+        def max_inner_product(self, other: Any) -> Any:
             return self.op('<#>', return_type=Float)(other)
 
-        def cosine_distance(self, other):
+        def cosine_distance(self, other: Any) -> Any:
             return self.op('<=>', return_type=Float)(other)
 
-        def l1_distance(self, other):
+        def l1_distance(self, other: Any) -> Any:
             return self.op('<+>', return_type=Float)(other)
 
 

pgvector/sqlalchemy/vector.py

@@ -1,49 +1,51 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
 from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy import Dialect
+from typing import Any
 from .. import Vector
 
 
 class VECTOR(UserDefinedType):
     cache_ok = True
     _string = String()
 
-    def __init__(self, dim=None):
+    def __init__(self, dim: int | None = None) -> None:
         super(UserDefinedType, self).__init__()
         self.dim = dim
 
-    def get_col_spec(self, **kw):
+    def get_col_spec(self, **kw) -> str:
         if self.dim is None:
             return 'VECTOR'
         return 'VECTOR(%d)' % self.dim
 
-    def bind_processor(self, dialect):
+    def bind_processor(self, dialect: Dialect) -> Any:
         def process(value):
             return Vector._to_db(value, self.dim)
         return process
 
-    def literal_processor(self, dialect):
+    def literal_processor(self, dialect: Dialect) -> Any:
         string_literal_processor = self._string._cached_literal_processor(dialect)
 
         def process(value):
             return string_literal_processor(Vector._to_db(value, self.dim))  # type: ignore
         return process
 
-    def result_processor(self, dialect, coltype):
+    def result_processor(self, dialect: Dialect, coltype: Any) -> Any:
         def process(value):
             return Vector._from_db(value)
         return process
 
     class comparator_factory(UserDefinedType.Comparator):
-        def l2_distance(self, other):
+        def l2_distance(self, other: Any) -> Any:
             return self.op('<->', return_type=Float)(other)
 
-        def max_inner_product(self, other):
+        def max_inner_product(self, other: Any) -> Any:
             return self.op('<#>', return_type=Float)(other)
 
-        def cosine_distance(self, other):
+        def cosine_distance(self, other: Any) -> Any:
             return self.op('<=>', return_type=Float)(other)
 
-        def l1_distance(self, other):
+        def l1_distance(self, other: Any) -> Any:
             return self.op('<+>', return_type=Float)(other)