Fixed typechecking for Python 3.10 [skip ci]

Author: ankane

pgvector/bit.py

@@ -9,7 +9,7 @@
 
 
 class Bit:
-    def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None:
+    def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]]) -> None:
         if isinstance(value, bytes):
             self._length = 8 * len(value)
             self._data = value
@@ -44,7 +44,7 @@ def bit_value(v: bool) -> str:
                 raise ValueError('expected ndim to be 1')
 
             self._length = len(value)
-            self._data = np.packbits(value).tobytes()
+            self._data = np.packbits(value).tobytes()  # type: ignore
         else:
             raise ValueError('expected bytes, str, list, or ndarray')
 
@@ -60,7 +60,7 @@ def to_list(self) -> list[bool]:
         # TODO improve
         return [v != '0' for v in self.to_text()]
 
-    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]:
+    def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]:
         return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool)
 
     def to_text(self) -> str:

pgvector/django/extensions.py

@@ -6,6 +6,6 @@
 class VectorExtension(CreateExtension):
     def __init__(self, hints: Any = None) -> None:
         if VERSION[0] >= 6:
-            super().__init__('vector', hints=hints)
+            super().__init__('vector', hints=hints)  # type: ignore
         else:
             self.name = 'vector'

pgvector/halfvec.py

@@ -11,7 +11,7 @@
 
 
 class HalfVector:
-    def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
+    def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None:
         if isinstance(value, list):
             dim = len(value)
             try:
@@ -43,7 +43,7 @@ def to_list(self) -> list[float]:
         dim = len(self._value)
         return list(struct.unpack(f'{dim}e', self._value))
 
-    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]:
+    def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float16]]:
         return np.frombuffer(self._value, dtype=np.float16)
 
     def to_text(self) -> str:

pgvector/psycopg/vector.py

@@ -18,7 +18,7 @@
 class VectorDumper(Dumper):
     format = Format.TEXT
 
-    def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None:
+    def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None:
         if not isinstance(obj, Vector):
             obj = Vector(obj)
         return obj.to_text().encode('utf8')
@@ -27,7 +27,7 @@ def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None:
 class VectorBinaryDumper(VectorDumper):
     format = Format.BINARY
 
-    def dump(self, obj: 'Vector | np.ndarray') -> Buffer | None:
+    def dump(self, obj: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> Buffer | None:
         if not isinstance(obj, Vector):
             obj = Vector(obj)
         return obj.to_binary()

pgvector/psycopg2/vector.py

@@ -10,7 +10,7 @@
 
 
 class VectorAdapter:
-    def __init__(self, value: 'Vector | np.ndarray') -> None:
+    def __init__(self, value: 'Vector | np.ndarray[tuple[int, ...], np.dtype[np.floating]]') -> None:
         if not isinstance(value, Vector):
             value = Vector(value)
         self._value = value

pgvector/sparsevec.py

@@ -22,10 +22,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] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | 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] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None:
         if SCIPY_AVAILABLE and isinstance(value, (sparray, spmatrix)):
             if dimensions is not NO_DEFAULT:
                 raise ValueError('extra argument')
@@ -70,7 +70,7 @@ def to_list(self) -> list[float]:
             vec[i] = v
         return vec
 
-    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]:
+    def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]:
         vec = np.repeat(0.0, self._dim).astype(np.float32)
         for i, v in zip(self._indices, self._values):
             vec[i] = v

pgvector/vector.py

@@ -11,7 +11,7 @@
 
 
 class Vector:
-    def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
+    def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]]) -> None:
         if isinstance(value, list):
             try:
                 self._value = array.array('f', value)
@@ -41,7 +41,7 @@ def dimensions(self) -> int:
     def to_list(self) -> list[float]:
         return self._value.tolist()
 
-    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]:
+    def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]:
         return np.frombuffer(self._value, dtype=np.float32)
 
     def to_text(self) -> str: