Added tests for dim methods [skip ci]

ankane · ankane · commit 87249e92ad2d · 2024-06-01T20:20:27.000-07:00
diff --git a/pgvector/utils/bit.py b/pgvector/utils/bit.py
@@ -27,8 +27,7 @@ def to_text(self):
         return ''.join(self._value.astype(np.uint8).astype(str))
 
     def to_binary(self):
-        value = self._value
-        return pack('>i', len(value)) + np.packbits(value).tobytes()
+        return pack('>i', len(self._value)) + np.packbits(self._value).tobytes()
 
     def from_text(value):
         return Bit(np.asarray([v != '0' for v in value], dtype=bool))
diff --git a/pgvector/utils/halfvec.py b/pgvector/utils/halfvec.py
@@ -17,7 +17,7 @@ def __repr__(self):
         return f'HalfVector({self.to_list()})'
 
     def dim(self):
-        return self._value.shape[0]
+        return len(self._value)
 
     def to_list(self):
         return self._value.tolist()
diff --git a/pgvector/utils/vector.py b/pgvector/utils/vector.py
@@ -17,7 +17,7 @@ def __repr__(self):
         return f'Vector({self.to_list()})'
 
     def dim(self):
-        return self._value.shape[0]
+        return len(self._value)
 
     def to_list(self):
         return self._value.tolist()
diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py
@@ -37,3 +37,6 @@ def test_ndim_zero(self):
     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_dim(self):
+        assert HalfVector([1, 2, 3]).dim() == 3
diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py
@@ -11,3 +11,6 @@ def test_from_dense(self):
     def test_repr(self):
         assert repr(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])'
         assert str(SparseVector.from_dense([1, 2, 3])) == 'SparseVector(3, [0, 1, 2], [1.0, 2.0, 3.0])'
+
+    def test_dim(self):
+        assert SparseVector.from_dense([1, 2, 3]).dim() == 3
diff --git a/tests/test_vector.py b/tests/test_vector.py
@@ -37,3 +37,6 @@ def test_ndim_zero(self):
     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_dim(self):
+        assert Vector([1, 2, 3]).dim() == 3
