Enhances DjangoAdmin support (pgvector#22)

jerinpetergeorge · ankane · commit 557983b36b14 · 2023-07-31T21:06:25.000-07:00
diff --git a/pgvector/django/__init__.py b/pgvector/django/__init__.py
@@ -1,6 +1,7 @@
 from django.contrib.postgres.operations import CreateExtension
 from django.contrib.postgres.indexes import PostgresIndex
 from django.db.models import Field, FloatField, Func, Value
+from .forms import VectorFormField
 from ..utils import from_db, to_db
 
 __all__ = ['VectorExtension', 'VectorField', 'IvfflatIndex', 'L2Distance', 'MaxInnerProduct', 'CosineDistance']
@@ -48,6 +49,9 @@ def validate(self, value, model_instance):
     def run_validators(self, value):
         super().run_validators(value.tolist())
 
+    def formfield(self, **kwargs):
+        return super().formfield(form_class=VectorFormField, **kwargs)
+
 
 class IvfflatIndex(PostgresIndex):
     suffix = 'ivfflat'
diff --git a/pgvector/django/forms.py b/pgvector/django/forms.py
@@ -0,0 +1,14 @@
+from django import forms
+from .widgets import VectorWidget
+
+
+class VectorFormField(forms.CharField):
+    widget = VectorWidget
+
+    def has_changed(self, initial, data):
+        try:
+            initial = initial.tolist()
+        except AttributeError:
+            # initial could be None
+            pass
+        return super().has_changed(initial, data)
diff --git a/pgvector/django/widgets.py b/pgvector/django/widgets.py
@@ -0,0 +1,11 @@
+from django import forms
+
+
+class VectorWidget(forms.TextInput):
+    def format_value(self, value):
+        try:
+            value = value.tolist()
+        except AttributeError:
+            # value could be None
+            pass
+        return super().format_value(value)
diff --git a/tests/test_django.py b/tests/test_django.py
@@ -136,10 +136,16 @@ def test_serialization(self):
                     obj.save()
 
     def test_form(self):
-        form = ItemForm(data={'embedding': [1, 2, 3]})
+        form = ItemForm(data={'embedding': "[1, 2, 3]"})
         assert form.is_valid()
         assert 'value="[1, 2, 3]"' in form.as_div()
 
+    def test_form_has_changed(self):
+        Item(id=1, embedding=[1, 2, 3]).save()
+        item = Item.objects.get(pk=1)
+        form = ItemForm(instance=item, data={'embedding': "[1, 2, 4]"})
+        assert form.has_changed()
+
     def test_form_instance(self):
         Item(id=1, embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
@@ -150,7 +156,7 @@ def test_form_instance(self):
     def test_form_save(self):
         Item(id=1, embedding=[1, 2, 3]).save()
         item = Item.objects.get(pk=1)
-        form = ItemForm(instance=item, data={'embedding': [4, 5, 6]})
+        form = ItemForm(instance=item, data={'embedding': "[4, 5, 6]"})
         assert form.is_valid()
         assert form.save()
-        assert [4, 5, 6], Item.objects.get(pk=1).embedding
+        assert [4, 5, 6] == Item.objects.get(pk=1).embedding.tolist()
