Improved type hints [skip ci]

Author: ankane

pgvector/halfvec.py

@@ -4,7 +4,7 @@
 
 
 class HalfVector:
-    def __init__(self, value: object) -> None:
+    def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None:
         # asarray still copies if same dtype
         if not isinstance(value, np.ndarray) or value.dtype != '>f2':
             value = np.asarray(value, dtype='>f2')
@@ -55,7 +55,7 @@ def _to_db(cls, value: object, dim: int | None = None) -> str | None:
             return value
 
         if not isinstance(value, cls):
-            value = cls(value)
+            value = cls(value)  # ty: ignore[invalid-argument-type]
 
         if dim is not None and value.dimensions() != dim:
             raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions()))
@@ -68,7 +68,7 @@ def _to_db_binary(cls, value: object) -> bytes | None:
             return value
 
         if not isinstance(value, cls):
-            value = cls(value)
+            value = cls(value)  # ty: ignore[invalid-argument-type]
 
         return value.to_binary()
 

pgvector/vector.py

@@ -4,7 +4,7 @@
 
 
 class Vector:
-    def __init__(self, value: object) -> None:
+    def __init__(self, value: list[float] | tuple[float, ...] | np.ndarray) -> None:
         # asarray still copies if same dtype
         if not isinstance(value, np.ndarray) or value.dtype != '>f4':
             value = np.asarray(value, dtype='>f4')
@@ -55,7 +55,7 @@ def _to_db(cls, value: object, dim: int | None = None) -> str | None:
             return value
 
         if not isinstance(value, cls):
-            value = cls(value)
+            value = cls(value)  # ty: ignore[invalid-argument-type]
 
         if dim is not None and value.dimensions() != dim:
             raise ValueError('expected %d dimensions, not %d' % (dim, value.dimensions()))
@@ -68,7 +68,7 @@ def _to_db_binary(cls, value: object) -> bytes | None:
             return value
 
         if not isinstance(value, cls):
-            value = cls(value)
+            value = cls(value)  # ty: ignore[invalid-argument-type]
 
         return value.to_binary()
 

tests/test_half_vector.py

@@ -10,7 +10,7 @@ def test_list(self):
 
     def test_list_str(self):
         with pytest.raises(ValueError, match='could not convert string to float'):
-            HalfVector([1, 'two', 3])
+            HalfVector([1, 'two', 3])  # ty: ignore[invalid-argument-type]
 
     def test_tuple(self):
         assert HalfVector((1, 2, 3)).to_list() == [1, 2, 3]
@@ -27,12 +27,12 @@ def test_ndarray_same_object(self):
 
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
-            HalfVector([[1, 2], [3, 4]])
+            HalfVector([[1, 2], [3, 4]])  # ty: ignore[invalid-argument-type]
         assert str(error.value) == 'expected ndim to be 1'
 
     def test_ndim_zero(self):
         with pytest.raises(ValueError) as error:
-            HalfVector(1)
+            HalfVector(1)  # ty: ignore[invalid-argument-type]
         assert str(error.value) == 'expected ndim to be 1'
 
     def test_repr(self):

tests/test_sqlalchemy.py

@@ -601,7 +601,7 @@ async def test_halfvec(self, engine):
                 embedding = [1, 2, 3]
                 session.add(Item(id=1, half_embedding=embedding))
                 item = await session.get_one(Item, 1)
-                assert item.half_embedding == HalfVector(embedding)
+                assert item.half_embedding == HalfVector(embedding)  # ty: ignore[invalid-argument-type]
 
         await engine.dispose()
 

tests/test_vector.py

@@ -10,7 +10,7 @@ def test_list(self):
 
     def test_list_str(self):
         with pytest.raises(ValueError, match='could not convert string to float'):
-            Vector([1, 'two', 3])
+            Vector([1, 'two', 3])  # ty: ignore[invalid-argument-type]
 
     def test_tuple(self):
         assert Vector((1, 2, 3)).to_list() == [1, 2, 3]
@@ -27,12 +27,12 @@ def test_ndarray_same_object(self):
 
     def test_ndim_two(self):
         with pytest.raises(ValueError) as error:
-            Vector([[1, 2], [3, 4]])
+            Vector([[1, 2], [3, 4]])  # ty: ignore[invalid-argument-type]
         assert str(error.value) == 'expected ndim to be 1'
 
     def test_ndim_zero(self):
         with pytest.raises(ValueError) as error:
-            Vector(1)
+            Vector(1)  # ty: ignore[invalid-argument-type]
         assert str(error.value) == 'expected ndim to be 1'
 
     def test_repr(self):