Improved typing [skip ci]

Author: ankane

pgvector/halfvec.py

@@ -86,16 +86,16 @@ def _from_text(cls, value: str) -> list[float]:
         return [float(v) for v in value[1:-1].split(',')]
 
     @classmethod
-    def _to_db(cls, value: object) -> str | None:
+    def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | HalfVector | None) -> str | None:
         if value is None:
             return value
 
         # fewer allocations for lists
         if isinstance(value, list):
             return f'[{",".join([str(float(v)) for v in value])}]'  # type: ignore
 
-        if not isinstance(value, cls):
-            value = cls(value)  # type: ignore
+        if not isinstance(value, HalfVector):
+            value = cls(value)
 
         return value.to_text()
 

pgvector/peewee/halfvec.py

@@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) ->
     def get_modifiers(self) -> list[int] | None:
         return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value: object) -> str | None:
+    def db_value(self, value: Any) -> str | None:
         return HalfVector._to_db(value)
 
     def python_value(self, value: Any) -> list[float] | None:

pgvector/peewee/sparsevec.py

@@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) ->
     def get_modifiers(self) -> list[int] | None:
         return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value: object) -> str | None:
+    def db_value(self, value: Any) -> str | None:
         return SparseVector._to_db(value)
 
     def python_value(self, value: Any) -> SparseVector | None:

pgvector/peewee/vector.py

@@ -13,7 +13,7 @@ def __init__(self, dimensions: int | None = None, *args: Any, **kwargs: Any) ->
     def get_modifiers(self) -> list[int] | None:
         return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value: object) -> str | None:
+    def db_value(self, value: Any) -> str | None:
         return Vector._to_db(value)
 
     def python_value(self, value: Any) -> list[float] | None:

pgvector/sparsevec.py

@@ -149,11 +149,11 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars
         return vec
 
     @classmethod
-    def _to_db(cls, value: object) -> str | None:
+    def _to_db(cls, value: list[float] | Any | SparseVector | None) -> str | None:
         if value is None:
             return value
 
-        if not isinstance(value, cls):
+        if not isinstance(value, SparseVector):
             value = cls(value)
 
         return value.to_text()

pgvector/vector.py

@@ -84,16 +84,16 @@ def _from_text(cls, value: str) -> list[float]:
         return [float(v) for v in value[1:-1].split(',')]
 
     @classmethod
-    def _to_db(cls, value: object) -> str | None:
+    def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | Vector | None) -> str | None:
         if value is None:
             return value
 
         # fewer allocations for lists
         if isinstance(value, list):
             return f'[{",".join([str(float(v)) for v in value])}]'  # type: ignore
 
-        if not isinstance(value, cls):
-            value = cls(value)  # type: ignore
+        if not isinstance(value, Vector):
+            value = cls(value)
 
         return value.to_text()