Vector: Fix type checking

crate · Dec 21, 2023 · 264f370 · 264f370
1 parent 82a9734
commit 264f370
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/src/sqlalchemy_cratedb/type/object.py b/src/sqlalchemy_cratedb/type/object.py
@@ -70,10 +70,8 @@ class ObjectTypeImpl(sqltypes.UserDefinedType, sqltypes.JSON):
     none_as_null = False
 
 
-ObjectType = MutableDict.as_mutable(ObjectTypeImpl)
-
-
 # Designated name to refer to. `Object` is too ambiguous.
+ObjectType = MutableDict.as_mutable(ObjectTypeImpl)
 
 # Backward-compatibility aliases.
 _deprecated_Craty = ObjectType
@@ -89,3 +87,6 @@ def __getattr__(name):
                       f"Please use ObjectType instead.", DeprecationWarning)
         return globals()[f"_deprecated_{name}"]
     raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+__all__ = deprecated_names
diff --git a/src/sqlalchemy_cratedb/type/vector.py b/src/sqlalchemy_cratedb/type/vector.py
@@ -44,7 +44,7 @@
 __all__ = ["FloatVector"]
 
 
-def from_db(value: t.Iterable) -> t.Optional[npt.ArrayLike]:
+def from_db(value: t.Iterable) -> t.Optional["npt.ArrayLike"]:
     import numpy as np
 
     # from `pgvector.utils`
@@ -146,13 +146,13 @@ def __init__(self, dimensions: int = None):
     def as_generic(self):
         return sa.ARRAY
 
-    def bind_processor(self, dialect: sa.Dialect) -> t.Callable:
+    def bind_processor(self, dialect: sa.engine.Dialect) -> t.Callable:
         def process(value: t.Iterable) -> t.Optional[t.List]:
             return to_db(value, self.dimensions)
 
         return process
 
-    def result_processor(self, dialect: sa.Dialect, coltype: t.Any) -> t.Callable:
+    def result_processor(self, dialect: sa.engine.Dialect, coltype: t.Any) -> t.Callable:
         def process(value: t.Any) -> t.Optional[npt.ArrayLike]:
             return from_db(value)