Improved type hints [skip ci]

Author: ankane

pgvector/peewee/bit.py

@@ -1,21 +1,22 @@
 from peewee import Expression, Field
+from typing import Any
 
 
 class FixedBitField(Field):
     field_type = 'bit'
 
-    def __init__(self, max_length=None, *args, **kwargs):
+    def __init__(self, max_length: int | None = None, *args, **kwargs) -> None:
         self.max_length = max_length
         super(FixedBitField, self).__init__(*args, **kwargs)
 
-    def get_modifiers(self):
-        return self.max_length and [self.max_length] or None
+    def get_modifiers(self) -> list[int] | None:
+        return [self.max_length] if self.max_length else None
 
-    def _distance(self, op, vector):
+    def _distance(self, op: str, vector: Any) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))
 
-    def hamming_distance(self, vector):
+    def hamming_distance(self, vector: Any) -> Expression:
         return self._distance('<~>', vector)
 
-    def jaccard_distance(self, vector):
+    def jaccard_distance(self, vector: Any) -> Expression:
         return self._distance('<%%>', vector)

pgvector/peewee/halfvec.py

@@ -1,34 +1,35 @@
 from peewee import Expression, Field
+from typing import Any
 from .. import HalfVector
 
 
 class HalfVectorField(Field):
     field_type = 'halfvec'
 
-    def __init__(self, dimensions=None, *args, **kwargs):
+    def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None:
         self.dimensions = dimensions
         super(HalfVectorField, self).__init__(*args, **kwargs)
 
-    def get_modifiers(self):
-        return self.dimensions and [self.dimensions] or None
+    def get_modifiers(self) -> list[int] | None:
+        return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value):
+    def db_value(self, value: Any) -> str | None:
         return HalfVector._to_db(value)
 
-    def python_value(self, value):
+    def python_value(self, value: Any) -> HalfVector | None:
         return HalfVector._from_db(value)
 
-    def _distance(self, op, vector):
+    def _distance(self, op: str, vector: Any) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))
 
-    def l2_distance(self, vector):
+    def l2_distance(self, vector: Any) -> Expression:
         return self._distance('<->', vector)
 
-    def max_inner_product(self, vector):
+    def max_inner_product(self, vector: Any) -> Expression:
         return self._distance('<#>', vector)
 
-    def cosine_distance(self, vector):
+    def cosine_distance(self, vector: Any) -> Expression:
         return self._distance('<=>', vector)
 
-    def l1_distance(self, vector):
+    def l1_distance(self, vector: Any) -> Expression:
         return self._distance('<+>', vector)

pgvector/peewee/sparsevec.py

@@ -1,34 +1,35 @@
 from peewee import Expression, Field
+from typing import Any
 from .. import SparseVector
 
 
 class SparseVectorField(Field):
     field_type = 'sparsevec'
 
-    def __init__(self, dimensions=None, *args, **kwargs):
+    def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None:
         self.dimensions = dimensions
         super(SparseVectorField, self).__init__(*args, **kwargs)
 
-    def get_modifiers(self):
-        return self.dimensions and [self.dimensions] or None
+    def get_modifiers(self) -> list[int] | None:
+        return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value):
+    def db_value(self, value: Any) -> str | None:
         return SparseVector._to_db(value)
 
-    def python_value(self, value):
+    def python_value(self, value: Any) -> SparseVector | None:
         return SparseVector._from_db(value)
 
-    def _distance(self, op, vector):
+    def _distance(self, op: str, vector: Any) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))
 
-    def l2_distance(self, vector):
+    def l2_distance(self, vector: Any) -> Expression:
         return self._distance('<->', vector)
 
-    def max_inner_product(self, vector):
+    def max_inner_product(self, vector: Any) -> Expression:
         return self._distance('<#>', vector)
 
-    def cosine_distance(self, vector):
+    def cosine_distance(self, vector: Any) -> Expression:
         return self._distance('<=>', vector)
 
-    def l1_distance(self, vector):
+    def l1_distance(self, vector: Any) -> Expression:
         return self._distance('<+>', vector)

pgvector/peewee/vector.py

@@ -1,34 +1,36 @@
+import numpy as np
 from peewee import Expression, Field
+from typing import Any
 from .. import Vector
 
 
 class VectorField(Field):
     field_type = 'vector'
 
-    def __init__(self, dimensions=None, *args, **kwargs):
+    def __init__(self, dimensions: int | None = None, *args, **kwargs) -> None:
         self.dimensions = dimensions
         super(VectorField, self).__init__(*args, **kwargs)
 
-    def get_modifiers(self):
-        return self.dimensions and [self.dimensions] or None
+    def get_modifiers(self) -> list[int] | None:
+        return [self.dimensions] if self.dimensions else None
 
-    def db_value(self, value):
+    def db_value(self, value: Any) -> str | None:
         return Vector._to_db(value)
 
-    def python_value(self, value):
+    def python_value(self, value: Any) -> np.ndarray | None:
         return Vector._from_db(value)
 
-    def _distance(self, op, vector):
+    def _distance(self, op: str, vector: Any) -> Expression:
         return Expression(lhs=self, op=op, rhs=self.to_value(vector))
 
-    def l2_distance(self, vector):
+    def l2_distance(self, vector: Any) -> Expression:
         return self._distance('<->', vector)
 
-    def max_inner_product(self, vector):
+    def max_inner_product(self, vector: Any) -> Expression:
         return self._distance('<#>', vector)
 
-    def cosine_distance(self, vector):
+    def cosine_distance(self, vector: Any) -> Expression:
         return self._distance('<=>', vector)
 
-    def l1_distance(self, vector):
+    def l1_distance(self, vector: Any) -> Expression:
         return self._distance('<+>', vector)

pgvector/sqlalchemy/halfvec.py

@@ -18,7 +18,7 @@ def get_col_spec(self, **kw) -> str:
             return 'HALFVEC'
         return 'HALFVEC(%d)' % self.dim
 
-    def bind_processor(self, dialect: Dialect):
+    def bind_processor(self, dialect: Dialect) -> Any:
         def process(value):
             return HalfVector._to_db(value, self.dim)
         return process