Restore 0.1.5 version from stash

This commit is contained in:
liaibo
2025-12-08 19:56:24 +08:00
parent de189e938d
commit 8db3f4e32d
8578 changed files with 2703426 additions and 217 deletions
@@ -0,0 +1,11 @@
from .register import register_vector, register_vector_async
from ..utils import Bit, HalfVector, SparseVector, Vector
__all__ = [
'register_vector',
'register_vector_async',
'Vector',
'HalfVector',
'Bit',
'SparseVector'
]
@@ -0,0 +1,31 @@
from psycopg.adapt import Dumper
from psycopg.pq import Format
from ..utils import Bit
class BitDumper(Dumper):
format = Format.TEXT
def dump(self, obj):
return Bit._to_db(obj).encode('utf8')
class BitBinaryDumper(BitDumper):
format = Format.BINARY
def dump(self, obj):
return Bit._to_db_binary(obj)
def register_bit_info(context, info):
info.register(context)
# add oid to anonymous class for set_types
text_dumper = type('', (BitDumper,), {'oid': info.oid})
binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid})
adapters = context.adapters
adapters.register_dumper(Bit, text_dumper)
adapters.register_dumper(Bit, binary_dumper)
@@ -0,0 +1,53 @@
from psycopg.adapt import Loader, Dumper
from psycopg.pq import Format
from ..utils import HalfVector
class HalfVectorDumper(Dumper):
format = Format.TEXT
def dump(self, obj):
return HalfVector._to_db(obj).encode('utf8')
class HalfVectorBinaryDumper(HalfVectorDumper):
format = Format.BINARY
def dump(self, obj):
return HalfVector._to_db_binary(obj)
class HalfVectorLoader(Loader):
format = Format.TEXT
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return HalfVector._from_db(data.decode('utf8'))
class HalfVectorBinaryLoader(HalfVectorLoader):
format = Format.BINARY
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return HalfVector._from_db_binary(data)
def register_halfvec_info(context, info):
info.register(context)
# add oid to anonymous class for set_types
text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid})
binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid})
adapters = context.adapters
adapters.register_dumper(HalfVector, text_dumper)
adapters.register_dumper(HalfVector, binary_dumper)
adapters.register_loader(info.oid, HalfVectorLoader)
adapters.register_loader(info.oid, HalfVectorBinaryLoader)
@@ -0,0 +1,37 @@
from psycopg.types import TypeInfo
from .bit import register_bit_info
from .halfvec import register_halfvec_info
from .sparsevec import register_sparsevec_info
from .vector import register_vector_info
def register_vector(context):
info = TypeInfo.fetch(context, 'vector')
register_vector_info(context, info)
info = TypeInfo.fetch(context, 'bit')
register_bit_info(context, info)
info = TypeInfo.fetch(context, 'halfvec')
if info is not None:
register_halfvec_info(context, info)
info = TypeInfo.fetch(context, 'sparsevec')
if info is not None:
register_sparsevec_info(context, info)
async def register_vector_async(context):
info = await TypeInfo.fetch(context, 'vector')
register_vector_info(context, info)
info = await TypeInfo.fetch(context, 'bit')
register_bit_info(context, info)
info = await TypeInfo.fetch(context, 'halfvec')
if info is not None:
register_halfvec_info(context, info)
info = await TypeInfo.fetch(context, 'sparsevec')
if info is not None:
register_sparsevec_info(context, info)
@@ -0,0 +1,53 @@
from psycopg.adapt import Loader, Dumper
from psycopg.pq import Format
from ..utils import SparseVector
class SparseVectorDumper(Dumper):
format = Format.TEXT
def dump(self, obj):
return SparseVector._to_db(obj).encode('utf8')
class SparseVectorBinaryDumper(SparseVectorDumper):
format = Format.BINARY
def dump(self, obj):
return SparseVector._to_db_binary(obj)
class SparseVectorLoader(Loader):
format = Format.TEXT
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return SparseVector._from_db(data.decode('utf8'))
class SparseVectorBinaryLoader(SparseVectorLoader):
format = Format.BINARY
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return SparseVector._from_db_binary(data)
def register_sparsevec_info(context, info):
info.register(context)
# add oid to anonymous class for set_types
text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid})
binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid})
adapters = context.adapters
adapters.register_dumper(SparseVector, text_dumper)
adapters.register_dumper(SparseVector, binary_dumper)
adapters.register_loader(info.oid, SparseVectorLoader)
adapters.register_loader(info.oid, SparseVectorBinaryLoader)
@@ -0,0 +1,58 @@
import psycopg
from psycopg.adapt import Loader, Dumper
from psycopg.pq import Format
from ..utils import Vector
class VectorDumper(Dumper):
format = Format.TEXT
def dump(self, obj):
return Vector._to_db(obj).encode('utf8')
class VectorBinaryDumper(VectorDumper):
format = Format.BINARY
def dump(self, obj):
return Vector._to_db_binary(obj)
class VectorLoader(Loader):
format = Format.TEXT
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return Vector._from_db(data.decode('utf8'))
class VectorBinaryLoader(VectorLoader):
format = Format.BINARY
def load(self, data):
if isinstance(data, memoryview):
data = bytes(data)
return Vector._from_db_binary(data)
def register_vector_info(context, info):
if info is None:
raise psycopg.ProgrammingError('vector type not found in the database')
info.register(context)
# add oid to anonymous class for set_types
text_dumper = type('', (VectorDumper,), {'oid': info.oid})
binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid})
adapters = context.adapters
adapters.register_dumper('numpy.ndarray', text_dumper)
adapters.register_dumper('numpy.ndarray', binary_dumper)
adapters.register_dumper(Vector, text_dumper)
adapters.register_dumper(Vector, binary_dumper)
adapters.register_loader(info.oid, VectorLoader)
adapters.register_loader(info.oid, VectorBinaryLoader)