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 HalfVector

class HalfVectorDumper(Dumper):

format = Format.TEXT

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

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

class HalfVectorBinaryDumper(HalfVectorDumper):

format = Format.BINARY

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

return obj.to_binary()

class HalfVectorLoader(Loader):

format = Format.TEXT

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

if isinstance(data, memoryview):

data = data.tobytes()

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

class HalfVectorBinaryLoader(HalfVectorLoader):

format = Format.BINARY

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

return HalfVector.from_binary(data)

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

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)