Simplified code

Author: ankane

pgvector/asyncpg/register.py

@@ -7,7 +7,7 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None:
         'vector',
         schema=schema,
         encoder=Vector._to_db_binary,
-        decoder=Vector._from_db_binary,
+        decoder=Vector.from_binary,
         format='binary'
     )
 
@@ -16,15 +16,15 @@ async def register_vector(conn: Connection, schema: str = 'public') -> None:
             'halfvec',
             schema=schema,
             encoder=HalfVector._to_db_binary,
-            decoder=HalfVector._from_db_binary,
+            decoder=HalfVector.from_binary,
             format='binary'
         )
 
         await conn.set_type_codec(
             'sparsevec',
             schema=schema,
             encoder=SparseVector._to_db_binary,
-            decoder=SparseVector._from_db_binary,
+            decoder=SparseVector.from_binary,
             format='binary'
         )
     except ValueError as e:

pgvector/django/halfvec.py

@@ -31,7 +31,7 @@ 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)
+            return HalfVector.from_text(value)
         else:
             return HalfVector(value)
 

pgvector/django/vector.py

@@ -31,7 +31,7 @@ def to_python(self, value: Any) -> Vector | None:
         if value is None or isinstance(value, Vector):
             return value
         elif isinstance(value, str):
-            return Vector._from_db(value)
+            return Vector.from_text(value)
         else:
             return Vector(value)
 

pgvector/halfvec.py

@@ -96,10 +96,3 @@ def _from_db(cls, value: str | HalfVector | None) -> HalfVector | None:
             return value
 
         return cls.from_text(value)
-
-    @classmethod
-    def _from_db_binary(cls, value: bytes | HalfVector | None) -> HalfVector | None:
-        if value is None or isinstance(value, HalfVector):
-            return value
-
-        return cls.from_binary(value)

pgvector/pg8000/register.py

@@ -11,7 +11,7 @@ def register_vector(conn: Connection) -> None:
         raise RuntimeError('vector type not found in the database')
 
     conn.register_out_adapter(Vector, Vector._to_db)
-    conn.register_in_adapter(type_info['vector'], Vector._from_db)
+    conn.register_in_adapter(type_info['vector'], Vector.from_text)
 
     try:
         import numpy as np
@@ -21,8 +21,8 @@ def register_vector(conn: Connection) -> None:
 
     if 'halfvec' in type_info:
         conn.register_out_adapter(HalfVector, HalfVector._to_db)
-        conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db)
+        conn.register_in_adapter(type_info['halfvec'], HalfVector.from_text)
 
     if 'sparsevec' in type_info:
         conn.register_out_adapter(SparseVector, SparseVector._to_db)
-        conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db)
+        conn.register_in_adapter(type_info['sparsevec'], SparseVector.from_text)

pgvector/psycopg/bit.py

@@ -12,14 +12,14 @@ class BitDumper(Dumper):
     format = Format.TEXT
 
     def dump(self, obj: Bit) -> Buffer | None:
-        return Bit._to_db(obj).encode('utf8')
+        return obj.to_text().encode('utf8')
 
 
 class BitBinaryDumper(BitDumper):
     format = Format.BINARY
 
     def dump(self, obj: Bit) -> Buffer | None:
-        return Bit._to_db_binary(obj)
+        return obj.to_binary()
 
 
 def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None) -> None:

pgvector/psycopg/halfvec.py

@@ -12,15 +12,14 @@ class HalfVectorDumper(Dumper):
     format = Format.TEXT
 
     def dump(self, obj: HalfVector) -> Buffer | None:
-        value = HalfVector._to_db(obj)
-        return value if value is None else value.encode('utf8')
+        return obj.to_text().encode('utf8')
 
 
 class HalfVectorBinaryDumper(HalfVectorDumper):
     format = Format.BINARY
 
     def dump(self, obj: HalfVector) -> Buffer | None:
-        return HalfVector._to_db_binary(obj)
+        return obj.to_binary()
 
 
 class HalfVectorLoader(Loader):
@@ -29,7 +28,7 @@ class HalfVectorLoader(Loader):
     def load(self, data: Buffer) -> HalfVector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
-        return HalfVector._from_db(data.decode('utf8'))
+        return HalfVector.from_text(data.decode('utf8'))
 
 
 class HalfVectorBinaryLoader(HalfVectorLoader):
@@ -38,7 +37,7 @@ class HalfVectorBinaryLoader(HalfVectorLoader):
     def load(self, data: Buffer) -> HalfVector | None:
         if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
-        return HalfVector._from_db_binary(data)
+        return HalfVector.from_binary(data)
 
 
 def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo) -> None:

pgvector/psycopg/sparsevec.py

@@ -12,15 +12,14 @@ class SparseVectorDumper(Dumper):
     format = Format.TEXT
 
     def dump(self, obj: SparseVector) -> Buffer | None:
-        value = SparseVector._to_db(obj)
-        return value if value is None else value.encode('utf8')
+        return obj.to_text().encode('utf8')
 
 
 class SparseVectorBinaryDumper(SparseVectorDumper):
     format = Format.BINARY
 
     def dump(self, obj: SparseVector) -> Buffer | None:
-        return SparseVector._to_db_binary(obj)
+        return obj.to_binary()
 
 
 class SparseVectorLoader(Loader):
@@ -29,7 +28,7 @@ class SparseVectorLoader(Loader):
     def load(self, data: Buffer) -> SparseVector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
-        return SparseVector._from_db(data.decode('utf8'))
+        return SparseVector.from_text(data.decode('utf8'))
 
 
 class SparseVectorBinaryLoader(SparseVectorLoader):
@@ -38,7 +37,7 @@ class SparseVectorBinaryLoader(SparseVectorLoader):
     def load(self, data: Buffer) -> SparseVector | None:
         if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
-        return SparseVector._from_db_binary(data)
+        return SparseVector.from_binary(data)
 
 
 def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo) -> None:

pgvector/psycopg/vector.py

@@ -16,15 +16,18 @@ class VectorDumper(Dumper):
     format = Format.TEXT
 
     def dump(self, obj: Vector | np.ndarray) -> Buffer | None:
-        value = Vector._to_db(obj)
-        return value if value is None else value.encode('utf8')
+        if not isinstance(obj, Vector):
+            obj = Vector(obj)
+        return obj.to_text().encode('utf8')
 
 
 class VectorBinaryDumper(VectorDumper):
     format = Format.BINARY
 
     def dump(self, obj: Vector | np.ndarray) -> Buffer | None:
-        return Vector._to_db_binary(obj)
+        if not isinstance(obj, Vector):
+            obj = Vector(obj)
+        return obj.to_binary()
 
 
 class VectorLoader(Loader):
@@ -33,7 +36,7 @@ class VectorLoader(Loader):
     def load(self, data: Buffer) -> Vector | None:
         if isinstance(data, memoryview):
             data = bytes(data)
-        return Vector._from_db(data.decode('utf8'))
+        return Vector.from_text(data.decode('utf8'))
 
 
 class VectorBinaryLoader(VectorLoader):
@@ -42,7 +45,7 @@ class VectorBinaryLoader(VectorLoader):
     def load(self, data: Buffer) -> Vector | None:
         if isinstance(data, (bytearray, memoryview)):
             data = bytes(data)
-        return Vector._from_db_binary(data)
+        return Vector.from_binary(data)
 
 
 def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) -> None:

pgvector/sparsevec.py

@@ -174,10 +174,3 @@ def _from_db(cls, value: str | SparseVector | None) -> SparseVector | None:
             return value
 
         return cls.from_text(value)
-
-    @classmethod
-    def _from_db_binary(cls, value: bytes | SparseVector | None) -> SparseVector | None:
-        if value is None or isinstance(value, SparseVector):
-            return value
-
-        return cls.from_binary(value)