pgvector-python/pgvector/psycopg/sparsevec.py at master · pgvector/pgvector-python · GitHub

51 lines (35 loc) · 1.54 KB

from psycopg import BaseConnection

from psycopg.abc import Buffer

from psycopg.adapt import Loader, Dumper

from psycopg.pq import Format

from psycopg.types import TypeInfo

from typing import Any

from .. import SparseVector

class SparseVectorDumper(Dumper):

format = Format.TEXT

def dump(self, obj: SparseVector) -> Buffer | None:

return obj.to_text().encode('utf8')

class SparseVectorBinaryDumper(SparseVectorDumper):

format = Format.BINARY

def dump(self, obj: SparseVector) -> Buffer | None:

return obj.to_binary()

class SparseVectorLoader(Loader):

format = Format.TEXT

def load(self, data: Buffer) -> SparseVector | None:

if isinstance(data, memoryview):

data = data.tobytes()

return SparseVector.from_text(data.decode('utf8'))

class SparseVectorBinaryLoader(SparseVectorLoader):

format = Format.BINARY

def load(self, data: Buffer) -> SparseVector | None:

return SparseVector.from_binary(data)

def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> None:

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)

Expand file tree