Fixed typechecking

Author: ankane

pgvector/sqlalchemy/bit.py

@@ -1,5 +1,5 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float
+from sqlalchemy.types import UserDefinedType, TypeEngine, Float
 from sqlalchemy import Dialect, Operators
 from typing import Any
 
@@ -28,7 +28,7 @@ def process(value: Any) -> Any:
         else:
             return super().bind_processor(dialect)
 
-    class comparator_factory(UserDefinedType.Comparator):
+    class comparator_factory(TypeEngine.Comparator):
         def hamming_distance(self, other: object) -> Operators:
             return self.op('<~>', return_type=Float)(other)
 

pgvector/sqlalchemy/halfvec.py

@@ -1,5 +1,5 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String
 from sqlalchemy import Dialect, Operators
 from typing import Any
 from .. import HalfVector
@@ -35,7 +35,7 @@ def process(value: Any) -> HalfVector | None:
             return HalfVector._from_db(value)
         return process
 
-    class comparator_factory(UserDefinedType.Comparator):
+    class comparator_factory(TypeEngine.Comparator):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)
 

pgvector/sqlalchemy/sparsevec.py

@@ -1,5 +1,5 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String
 from sqlalchemy import Dialect, Operators
 from typing import Any
 from .. import SparseVector
@@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None:
             return SparseVector._from_db(value)
         return process
 
-    class comparator_factory(UserDefinedType.Comparator):
+    class comparator_factory(TypeEngine.Comparator):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)
 

pgvector/sqlalchemy/vector.py

@@ -1,5 +1,5 @@
 from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float, String
+from sqlalchemy.types import UserDefinedType, TypeEngine, Float, String
 from sqlalchemy import Dialect, Operators
 from typing import Any
 from .. import Vector
@@ -35,7 +35,7 @@ def process(value: Any) -> Vector | None:
             return Vector._from_db(value)
         return process
 
-    class comparator_factory(UserDefinedType.Comparator):
+    class comparator_factory(TypeEngine.Comparator):
         def l2_distance(self, other: object) -> Operators:
             return self.op('<->', return_type=Float)(other)