Added ndarray type alias

Author: ankane

pgvector/_utils.py

@@ -1,4 +1,10 @@
 import sys
+from typing import TYPE_CHECKING, TypeAlias
+
+if TYPE_CHECKING:
+    import numpy as np
+
+    ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]]
 
 
 def is_ndarray(value: object) -> bool:

pgvector/halfvec.py

@@ -7,10 +7,11 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    from ._utils import ndarray
 
 
 class HalfVector:
-    def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None:
+    def __init__(self, value: list[float] | ndarray) -> None:
         if isinstance(value, list):
             dim = len(value)
             try:
@@ -88,7 +89,7 @@ def _from_text(cls, value: str) -> list[float]:
         return [float(v) for v in value[1:-1].split(',')]
 
     @classmethod
-    def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | HalfVector | None) -> str | None:
+    def _to_db(cls, value: list[float] | ndarray | HalfVector | None) -> str | None:
         if value is None:
             return value
 

pgvector/psycopg/vector.py

@@ -10,13 +10,13 @@
 Buffer: TypeAlias = bytes | bytearray | memoryview
 
 if TYPE_CHECKING:
-    import numpy as np
+    from .._utils import ndarray
 
 
 class VectorDumper(Dumper):
     format = Format.TEXT
 
-    def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None:
+    def dump(self, obj: Vector | ndarray) -> Buffer | None:
         if not isinstance(obj, Vector):
             obj = Vector(obj)
         return obj.to_text().encode('utf8')
@@ -25,7 +25,7 @@ def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]])
 class VectorBinaryDumper(VectorDumper):
     format = Format.BINARY
 
-    def dump(self, obj: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> Buffer | None:
+    def dump(self, obj: Vector | ndarray) -> Buffer | None:
         if not isinstance(obj, Vector):
             obj = Vector(obj)
         return obj.to_binary()

pgvector/psycopg2/vector.py

@@ -4,11 +4,11 @@
 from .. import Vector
 
 if TYPE_CHECKING:
-    import numpy as np
+    from .._utils import ndarray
 
 
 class VectorAdapter:
-    def __init__(self, value: Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None:
+    def __init__(self, value: Vector | ndarray) -> None:
         if not isinstance(value, Vector):
             value = Vector(value)
         self._value = value

pgvector/sparsevec.py

@@ -6,6 +6,7 @@
 if TYPE_CHECKING:
     import numpy as np
     from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix
+    from ._utils import ndarray
 
 
 NO_DEFAULT = object()
@@ -17,10 +18,10 @@ def __init__(self, value: dict[int, float], dimensions: int, /) -> None:
         ...
 
     @overload
-    def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, /) -> None:
+    def __init__(self, value: list[float] | ndarray | sparray | spmatrix, /) -> None:
         ...
 
-    def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None:
+    def __init__(self, value: dict[int, float] | list[float] | ndarray | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None:
         if is_sparse_array(value):
             if dimensions is not NO_DEFAULT:
                 raise ValueError('extra argument')
@@ -107,7 +108,7 @@ def _from_sparse(self, arr: sparray | spmatrix) -> None:
             self._indices = value.col.tolist()
         self._values = value.data.tolist()
 
-    def _from_dense(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
+    def _from_dense(self, value: list[float] | ndarray) -> None:
         self._dim = len(value)
         self._indices = [i for i, v in enumerate(value) if v != 0]
         self._values = [float(value[i]) for i in self._indices]
@@ -148,7 +149,7 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float]) -> Spars
         return vec
 
     @classmethod
-    def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | sparray | spmatrix | SparseVector | None) -> str | None:
+    def _to_db(cls, value: list[float] | ndarray | sparray | spmatrix | SparseVector | None) -> str | None:
         if value is None:
             return value
 

pgvector/vector.py

@@ -7,10 +7,11 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    from ._utils import ndarray
 
 
 class Vector:
-    def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None:
+    def __init__(self, value: list[float] | ndarray) -> None:
         if isinstance(value, list):
             try:
                 self._value = array.array('f', value)
@@ -86,7 +87,7 @@ def _from_text(cls, value: str) -> list[float]:
         return [float(v) for v in value[1:-1].split(',')]
 
     @classmethod
-    def _to_db(cls, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]] | Vector | None) -> str | None:
+    def _to_db(cls, value: list[float] | ndarray | Vector | None) -> str | None:
         if value is None:
             return value