Improved typing [skip ci]

Author: ankane

pgvector/sqlalchemy/bit.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 
-class BIT(UserDefinedType):
+class BIT(UserDefinedType[Any]):
     cache_ok = True
 
     def __init__(self, length: int | None = None) -> None:
@@ -28,7 +28,7 @@ def process(value: Any) -> Any:
         else:
             return super().bind_processor(dialect)
 
-    class Comparator(TypeEngine.Comparator):
+    class Comparator(TypeEngine.Comparator[Any]):
         def hamming_distance(self, other: object) -> Operators:
             return self.op('<~>', return_type=Float)(other)
 

pgvector/sqlalchemy/halfvec.py

@@ -5,7 +5,7 @@
 from .. import HalfVector
 
 
-class HALFVEC(UserDefinedType):
+class HALFVEC(UserDefinedType[Any]):
     cache_ok = True
     _string = String()
 
@@ -35,7 +35,7 @@ def process(value: Any) -> HalfVector | None:
             return HalfVector._from_db(value)
         return process
 
-    class Comparator(TypeEngine.Comparator):
+    class Comparator(TypeEngine.Comparator[Any]):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)
 

pgvector/sqlalchemy/sparsevec.py

@@ -5,7 +5,7 @@
 from .. import SparseVector
 
 
-class SPARSEVEC(UserDefinedType):
+class SPARSEVEC(UserDefinedType[Any]):
     cache_ok = True
     _string = String()
 
@@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None:
             return SparseVector._from_db(value)
         return process
 
-    class Comparator(TypeEngine.Comparator):
+    class Comparator(TypeEngine.Comparator[Any]):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)
 

pgvector/sqlalchemy/vector.py

@@ -5,7 +5,7 @@
 from .. import Vector
 
 
-class VECTOR(UserDefinedType):
+class VECTOR(UserDefinedType[Any]):
     cache_ok = True
     _string = String()
 
@@ -35,7 +35,7 @@ def process(value: Any) -> Vector | None:
             return Vector._from_db(value)
         return process
 
-    class Comparator(TypeEngine.Comparator):
+    class Comparator(TypeEngine.Comparator[Any]):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)