Improved type hints [skip ci]

Author: ankane

pgvector/halfvec.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 import numpy as np
 from struct import pack, unpack_from
-from typing import Any
 
 
 class HalfVector:
-    def __init__(self, value: Any) -> None:
+    def __init__(self, value: object) -> None:
         # asarray still copies if same dtype
         if not isinstance(value, np.ndarray) or value.dtype != '>f2':
             value = np.asarray(value, dtype='>f2')
 
+        # for mypy
+        assert isinstance(value, np.ndarray)
+
         if value.ndim != 1:
             raise ValueError('expected ndim to be 1')
 
-        self._value = value
+        self._value = np.atleast_1d(value)
 
     def __repr__(self) -> str:
         return f'HalfVector({self.to_list()})'

pgvector/sparsevec.py

@@ -7,7 +7,7 @@
 
 
 class SparseVector:
-    def __init__(self, value: Any, dimensions: int | Any = NO_DEFAULT, /) -> None:
+    def __init__(self, value: dict[int, float] | list[float] | Any, dimensions: int | Any = NO_DEFAULT, /) -> None:
         if value.__class__.__module__.startswith('scipy.sparse.'):
             if dimensions is not NO_DEFAULT:
                 raise ValueError('extra argument')

pgvector/vector.py

@@ -1,19 +1,21 @@
 from __future__ import annotations
 import numpy as np
 from struct import pack, unpack_from
-from typing import Any
 
 
 class Vector:
-    def __init__(self, value: Any) -> None:
+    def __init__(self, value: object) -> None:
         # asarray still copies if same dtype
         if not isinstance(value, np.ndarray) or value.dtype != '>f4':
             value = np.asarray(value, dtype='>f4')
 
+        # for mypy
+        assert isinstance(value, np.ndarray)
+
         if value.ndim != 1:
             raise ValueError('expected ndim to be 1')
 
-        self._value = value
+        self._value = np.atleast_1d(value)
 
     def __repr__(self) -> str:
         return f'Vector({self.to_list()})'