Improved typing [skip ci]

Author: ankane

pyproject.toml

@@ -46,7 +46,7 @@ files = [
     "tests"
 ]
 strict = true
-disable_error_code = ["unused-ignore", "no-untyped-def", "no-untyped-call"]
+disable_error_code = ["unused-ignore", "no-untyped-def"]
 ignore_missing_imports = true
 disallow_subclassing_any = false
 python_executable = ".venv/bin/python"

tests/test_asyncpg.py

@@ -1,4 +1,5 @@
 import asyncpg
+from asyncpg import Connection
 from pgvector import HalfVector, SparseVector, Vector
 from pgvector.asyncpg import register_vector
 import pytest
@@ -11,7 +12,7 @@
 
 
 class TestAsyncpg:
-    async def setup_connection(self):
+    async def setup_connection(self) -> Connection:
         conn = await asyncpg.connect(database='pgvector_python_test')
         await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
         await register_vector(conn)

tests/test_peewee.py

@@ -31,7 +31,7 @@ class Meta:  # type: ignore
 db.create_tables([Item])
 
 
-def create_items():
+def create_items() -> None:
     Item.create(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1]))
     Item.create(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2]))
     Item.create(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2]))

tests/test_sqlalchemy.py

@@ -115,15 +115,15 @@ class Item(Base):
 binary_quantize_index.create(setup_engine)
 
 
-def create_items():
+def create_items() -> None:
     with Session(setup_engine) as session:
         session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])))
         session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])))
         session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])))
         session.commit()
 
 
-def delete_items():
+def delete_items() -> None:
     with Session(setup_engine) as session:
         session.query(Item).delete()
         session.commit()

tests/test_sqlmodel.py

@@ -39,7 +39,7 @@ class Item(SQLModel, table=True):
 index.create(engine)
 
 
-def create_items():
+def create_items() -> None:
     with Session(engine) as session:
         session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])))
         session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])))