Added more type hints [skip ci]

Author: ankane

pgvector/django/bit.py

@@ -1,32 +1,33 @@
 from django import forms
 from django.db.models import Field
+from typing import Any
 
 
 # https://docs.djangoproject.com/en/5.0/howto/custom-model-fields/
 class BitField(Field):
     description = 'Bit string'
 
-    def __init__(self, *args, length=None, **kwargs):
+    def __init__(self, *args: Any, length: int | None = None, **kwargs: Any) -> None:
         self.length = length
         super().__init__(*args, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         name, path, args, kwargs = super().deconstruct()
         if self.length is not None:
             kwargs['length'] = self.length
         return name, path, args, kwargs
 
-    def db_type(self, connection):
+    def db_type(self, connection: Any) -> str:
         if self.length is None:
             return 'bit'
         return 'bit(%d)' % self.length
 
-    def formfield(self, **kwargs):  # type: ignore
+    def formfield(self, **kwargs: Any):  # type: ignore
         return super().formfield(form_class=BitFormField, **kwargs)
 
 
 class BitFormField(forms.CharField):
-    def to_python(self, value):
+    def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None
         return super().to_python(value)

pgvector/django/extensions.py

@@ -1,9 +1,10 @@
 from django import VERSION
 from django.contrib.postgres.operations import CreateExtension
+from typing import Any
 
 
 class VectorExtension(CreateExtension):
-    def __init__(self, hints=None):
+    def __init__(self, hints: Any = None) -> None:
         if VERSION[0] >= 6:
             super().__init__('vector', hints=hints)
         else:

pgvector/django/functions.py

@@ -1,11 +1,12 @@
 from django.db.models import FloatField, Func, Value
 from .. import Vector, HalfVector, SparseVector
+from typing import Any
 
 
 class DistanceBase(Func):
     output_field = FloatField()
 
-    def __init__(self, expression, vector, **extra):
+    def __init__(self, expression: Any, vector: Any, **extra: Any) -> None:
         if not hasattr(vector, 'resolve_expression'):
             if isinstance(vector, HalfVector):
                 vector = Value(HalfVector._to_db(vector))
@@ -23,7 +24,7 @@ def __init__(self, expression, vector, **extra):
 class BitDistanceBase(Func):
     output_field = FloatField()
 
-    def __init__(self, expression, vector, **extra):
+    def __init__(self, expression: Any, vector: Any, **extra: Any) -> None:
         if not hasattr(vector, 'resolve_expression'):
             vector = Value(vector)
         super().__init__(expression, vector, **extra)

pgvector/django/halfvec.py

@@ -1,5 +1,6 @@
 from django import forms
 from django.db.models import Field
+from typing import Any
 from .. import HalfVector
 
 
@@ -8,44 +9,44 @@ class HalfVectorField(Field):
     description = 'Half vector'
     empty_strings_allowed = False
 
-    def __init__(self, *args, dimensions=None, **kwargs):
+    def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None:
         self.dimensions = dimensions
         super().__init__(*args, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         name, path, args, kwargs = super().deconstruct()
         if self.dimensions is not None:
             kwargs['dimensions'] = self.dimensions
         return name, path, args, kwargs
 
-    def db_type(self, connection):
+    def db_type(self, connection: Any) -> str:
         if self.dimensions is None:
             return 'halfvec'
         return 'halfvec(%d)' % self.dimensions
 
-    def from_db_value(self, value, expression, connection):
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> HalfVector | None:
         return HalfVector._from_db(value)
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> HalfVector | None:
         if value is None or isinstance(value, HalfVector):
             return value
         elif isinstance(value, str):
             return HalfVector._from_db(value)
         else:
             return HalfVector(value)
 
-    def get_prep_value(self, value):
+    def get_prep_value(self, value: Any) -> str | None:
         return HalfVector._to_db(value)
 
-    def value_to_string(self, obj):
+    def value_to_string(self, obj: Any) -> str | None:
         return self.get_prep_value(self.value_from_object(obj))
 
     def formfield(self, **kwargs):  # type: ignore
         return super().formfield(form_class=HalfVectorFormField, **kwargs)
 
 
 class HalfVectorWidget(forms.TextInput):
-    def format_value(self, value):
+    def format_value(self, value: Any) -> str | None:
         if isinstance(value, HalfVector):
             value = value.to_list()
         return super().format_value(value)
@@ -54,7 +55,7 @@ def format_value(self, value):
 class HalfVectorFormField(forms.CharField):
     widget = HalfVectorWidget
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None
         return super().to_python(value)

pgvector/django/indexes.py

@@ -1,20 +1,21 @@
 from django.contrib.postgres.indexes import PostgresIndex
+from typing import Any
 
 
 class IvfflatIndex(PostgresIndex):
     suffix = 'ivfflat'
 
-    def __init__(self, *expressions, lists=None, **kwargs):
+    def __init__(self, *expressions: Any, lists: int | None = None, **kwargs: Any) -> None:
         self.lists = lists
         super().__init__(*expressions, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         path, args, kwargs = super().deconstruct()
         if self.lists is not None:
             kwargs['lists'] = self.lists
         return path, args, kwargs
 
-    def get_with_params(self):
+    def get_with_params(self) -> list[str]:
         with_params = []
         if self.lists is not None:
             with_params.append('lists = %d' % self.lists)
@@ -24,20 +25,20 @@ def get_with_params(self):
 class HnswIndex(PostgresIndex):
     suffix = 'hnsw'
 
-    def __init__(self, *expressions, m=None, ef_construction=None, **kwargs):
+    def __init__(self, *expressions: Any, m: int | None = None, ef_construction: int | None = None, **kwargs: Any) -> None:
         self.m = m
         self.ef_construction = ef_construction
         super().__init__(*expressions, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         path, args, kwargs = super().deconstruct()
         if self.m is not None:
             kwargs['m'] = self.m
         if self.ef_construction is not None:
             kwargs['ef_construction'] = self.ef_construction
         return path, args, kwargs
 
-    def get_with_params(self):
+    def get_with_params(self) -> list[str]:
         with_params = []
         if self.m is not None:
             with_params.append('m = %d' % self.m)

pgvector/django/sparsevec.py

@@ -1,5 +1,6 @@
 from django import forms
 from django.db.models import Field
+from typing import Any
 from .. import SparseVector
 
 
@@ -8,39 +9,39 @@ class SparseVectorField(Field):
     description = 'Sparse vector'
     empty_strings_allowed = False
 
-    def __init__(self, *args, dimensions=None, **kwargs):
+    def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any):
         self.dimensions = dimensions
         super().__init__(*args, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         name, path, args, kwargs = super().deconstruct()
         if self.dimensions is not None:
             kwargs['dimensions'] = self.dimensions
         return name, path, args, kwargs
 
-    def db_type(self, connection):
+    def db_type(self, connection: Any) -> str:
         if self.dimensions is None:
             return 'sparsevec'
         return 'sparsevec(%d)' % self.dimensions
 
-    def from_db_value(self, value, expression, connection):
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> SparseVector | None:
         return SparseVector._from_db(value)
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> SparseVector | None:
         return SparseVector._from_db(value)
 
-    def get_prep_value(self, value):
+    def get_prep_value(self, value: Any) -> str | None:
         return SparseVector._to_db(value)
 
-    def value_to_string(self, obj):
+    def value_to_string(self, obj: Any) -> str | None:
         return self.get_prep_value(self.value_from_object(obj))
 
-    def formfield(self, **kwargs):  # type: ignore
+    def formfield(self, **kwargs: Any):  # type: ignore
         return super().formfield(form_class=SparseVectorFormField, **kwargs)
 
 
 class SparseVectorWidget(forms.TextInput):
-    def format_value(self, value):
+    def format_value(self, value: Any) -> str | None:
         if isinstance(value, SparseVector):
             value = value.to_text()
         return super().format_value(value)
@@ -49,7 +50,7 @@ def format_value(self, value):
 class SparseVectorFormField(forms.CharField):
     widget = SparseVectorWidget
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None
         return super().to_python(value)

pgvector/django/vector.py

@@ -1,6 +1,7 @@
 from django import forms
 from django.db.models import Field
 import numpy as np
+from typing import Any
 from .. import Vector
 
 
@@ -9,51 +10,51 @@ class VectorField(Field):
     description = 'Vector'
     empty_strings_allowed = False
 
-    def __init__(self, *args, dimensions=None, **kwargs):
+    def __init__(self, *args: Any, dimensions: int | None = None, **kwargs: Any) -> None:
         self.dimensions = dimensions
         super().__init__(*args, **kwargs)
 
-    def deconstruct(self):
+    def deconstruct(self) -> tuple:
         name, path, args, kwargs = super().deconstruct()
         if self.dimensions is not None:
             kwargs['dimensions'] = self.dimensions
         return name, path, args, kwargs
 
-    def db_type(self, connection):
+    def db_type(self, connection: Any) -> str:
         if self.dimensions is None:
             return 'vector'
         return 'vector(%d)' % self.dimensions
 
-    def from_db_value(self, value, expression, connection):
+    def from_db_value(self, value: Any, expression: Any, connection: Any) -> np.ndarray | None:
         return Vector._from_db(value)
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> np.ndarray | None:
         if isinstance(value, list):
             return np.array(value, dtype=np.float32)
         return Vector._from_db(value)
 
-    def get_prep_value(self, value):
+    def get_prep_value(self, value: Any) -> str | None:
         return Vector._to_db(value)
 
-    def value_to_string(self, obj):
+    def value_to_string(self, obj: Any) -> str | None:
         return self.get_prep_value(self.value_from_object(obj))
 
-    def validate(self, value, model_instance):
+    def validate(self, value: Any, model_instance: Any) -> None:
         if isinstance(value, np.ndarray):
             value = value.tolist()
         super().validate(value, model_instance)
 
-    def run_validators(self, value):
+    def run_validators(self, value: Any) -> None:
         if isinstance(value, np.ndarray):
             value = value.tolist()
         super().run_validators(value)
 
-    def formfield(self, **kwargs):  # type: ignore
+    def formfield(self, **kwargs: Any):  # type: ignore
         return super().formfield(form_class=VectorFormField, **kwargs)
 
 
 class VectorWidget(forms.TextInput):
-    def format_value(self, value):
+    def format_value(self, value: Any) -> str | None:
         if isinstance(value, np.ndarray):
             value = value.tolist()
         return super().format_value(value)
@@ -62,12 +63,12 @@ def format_value(self, value):
 class VectorFormField(forms.CharField):
     widget = VectorWidget
 
-    def has_changed(self, initial, data):
+    def has_changed(self, initial: Any, data: Any) -> bool:
         if isinstance(initial, np.ndarray):
             initial = initial.tolist()
         return super().has_changed(initial, data)
 
-    def to_python(self, value):
+    def to_python(self, value: Any) -> Any:
         if isinstance(value, str) and value == '':
             return None
         return super().to_python(value)