Improved type hints [skip ci]

Author: ankane

pgvector/bit.py

@@ -44,7 +44,7 @@ def __eq__(self, other: object) -> bool:
     def to_list(self) -> list[bool]:
         return self.to_numpy().tolist()
 
-    def to_numpy(self) -> np.ndarray:
+    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]:
         return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._len).astype(bool)
 
     def to_text(self) -> str:

pgvector/halfvec.py

@@ -32,7 +32,8 @@ def dimensions(self) -> int:
     def to_list(self) -> list[float]:
         return self._value.tolist()
 
-    def to_numpy(self) -> np.ndarray:
+    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float16]]:
+        # TODO return native endian
         return self._value
 
     def to_text(self) -> str:

pgvector/vector.py

@@ -32,7 +32,8 @@ def dimensions(self) -> int:
     def to_list(self) -> list[float]:
         return self._value.tolist()
 
-    def to_numpy(self) -> np.ndarray:
+    def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.float32]]:
+        # TODO return native endian
         return self._value
 
     def to_text(self) -> str: