Fixed equality for types [skip ci]

Author: ankane

pgvector/django/functions.py

@@ -13,6 +13,10 @@ def __init__(self, expression, vector, **extra):
                 vector = Value(SparseVector._to_db(vector))
             else:
                 vector = Value(Vector._to_db(vector))
+
+            # prevent error with unhashable types
+            self._constructor_args = ((expression, vector), extra)
+
         super().__init__(expression, vector, **extra)
 
 

pgvector/utils/bit.py

@@ -24,6 +24,11 @@ def __init__(self, value):
     def __repr__(self):
         return f'Bit({self.to_text()})'
 
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            return np.array_equal(self.to_numpy(), other.to_numpy())
+        return False
+
     def to_list(self):
         return self._value.tolist()
 

pgvector/utils/halfvec.py

@@ -16,6 +16,11 @@ def __init__(self, value):
     def __repr__(self):
         return f'HalfVector({self.to_list()})'
 
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            return np.array_equal(self.to_numpy(), other.to_numpy())
+        return False
+
     def dimensions(self):
         return len(self._value)
 

pgvector/utils/sparsevec.py

@@ -26,6 +26,11 @@ def __repr__(self):
         elements = dict(zip(self._indices, self._values))
         return f'SparseVector({elements}, {self._dim})'
 
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            return self.dimensions() == other.dimensions() and self.indices() == other.indices() and self.values() == other.values()
+        return False
+
     def dimensions(self):
         return self._dim
 

pgvector/utils/vector.py

@@ -16,6 +16,11 @@ def __init__(self, value):
     def __repr__(self):
         return f'Vector({self.to_list()})'
 
+    def __eq__(self, other):
+        if isinstance(other, self.__class__):
+            return np.array_equal(self.to_numpy(), other.to_numpy())
+        return False
+
     def dimensions(self):
         return len(self._value)
 

tests/test_bit.py

@@ -37,3 +37,7 @@ def test_ndim_zero(self):
     def test_repr(self):
         assert repr(Bit([True, False, True])) == 'Bit(101)'
         assert str(Bit([True, False, True])) == 'Bit(101)'
+
+    def test_equality(self):
+        assert Bit([True, False, True]) == Bit([True, False, True])
+        assert Bit([True, False, True]) != Bit([True, False, False])

tests/test_half_vector.py

@@ -38,5 +38,9 @@ def test_repr(self):
         assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])'
         assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])'
 
+    def test_equality(self):
+        assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3])
+        assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4])
+
     def test_dimensions(self):
         assert HalfVector([1, 2, 3]).dimensions() == 3

tests/test_sparse_vector.py

@@ -52,6 +52,11 @@ def test_repr(self):
         assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)'
         assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)'
 
+    def test_equality(self):
+        assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0])
+        assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1])
+        assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6)
+
     def test_dimensions(self):
         assert SparseVector([1, 0, 2, 0, 3, 0]).dimensions() == 6
 

tests/test_vector.py

@@ -38,5 +38,9 @@ def test_repr(self):
         assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])'
         assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])'
 
+    def test_equality(self):
+        assert Vector([1, 2, 3]) == Vector([1, 2, 3])
+        assert Vector([1, 2, 3]) != Vector([1, 2, 4])
+
     def test_dimensions(self):
         assert Vector([1, 2, 3]).dimensions() == 3