From a8df605bd02dfc6606d29a419bb883edbd668c23 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 3 Nov 2025 13:23:57 +0000 Subject: [PATCH] feat: add make_sequence parameter to register_dumper Allow to customize the way to extract a sequence from an object to dump it as composite type. --- docs/basic/pgtypes.rst | 4 +- docs/news.rst | 6 +- psycopg/psycopg/types/composite.py | 228 +++++++++++++++++++++-------- tests/types/test_composite.py | 75 +++++++--- 4 files changed, 230 insertions(+), 83 deletions(-) diff --git a/docs/basic/pgtypes.rst b/docs/basic/pgtypes.rst index 1d47feec4..cfca67cdc 100644 --- a/docs/basic/pgtypes.rst +++ b/docs/basic/pgtypes.rst @@ -66,7 +66,8 @@ using `~psycopg.types.composite.register_composite()`. .. versionadded:: 3.3 the `!make_instance` parameter -Example:: +Example: registering a composite without `!factory` information will create a +type on the fly, stored in `!CompositeInfo.python_type`. :: >>> from psycopg.types.composite import CompositeInfo, register_composite @@ -119,6 +120,7 @@ the composite to the right type requirements. For example:: >>> conn.execute("select '(1,spades)'::card").fetchone()[0] Card(suit='spades', value=1) + .. index:: pair: range; Data types diff --git a/docs/news.rst b/docs/news.rst index 337972851..4ad615743 100644 --- a/docs/news.rst +++ b/docs/news.rst @@ -16,9 +16,9 @@ Psycopg 3.3.0 (unreleased) .. rubric:: New top-level features - Add :ref:`template strings queries ` (:ticket:`#1054`). -- Add `!make_instance` parameter to `~types.composite.register_composite()` - to load objects from composite via keyword arguments or in other ways other - than just by values sequence (:ticket:`#932`). +- More flexible :ref:`composite adaptation`: it is now possible + to adapt Python objects to PostgreSQL composites and back even if they are not + sequences or if they take keyword arguments (:ticket:`#932`, :ticket:`#1202`). - Cursors are now iterators, not only iterables. This means you can call ``next(cur)`` to fetch the next row (:ticket:`#1064`). - Add `Cursor.results()` to iterate over the result sets of the queries diff --git a/psycopg/psycopg/types/composite.py b/psycopg/psycopg/types/composite.py index 520b48d92..a8adbab0b 100644 --- a/psycopg/psycopg/types/composite.py +++ b/psycopg/psycopg/types/composite.py @@ -8,8 +8,9 @@ from __future__ import annotations import re import struct +import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeVar, cast +from typing import TYPE_CHECKING, Any, Generic, NamedTuple, TypeAlias, TypeVar, cast from functools import cache from collections import namedtuple from collections.abc import Callable, Sequence @@ -25,6 +26,8 @@ from .._encodings import _as_python_identifier if TYPE_CHECKING: from .._connection_base import BaseConnection +logger = logging.getLogger("psycopg") + _struct_oidlen = struct.Struct("!Ii") _pack_oidlen = cast(Callable[[int, int], bytes], _struct_oidlen.pack) _unpack_oidlen = cast( @@ -32,6 +35,8 @@ _unpack_oidlen = cast( ) T = TypeVar("T") +InstanceMaker: TypeAlias = Callable[[Sequence[Any], Sequence[str]], T] +SequenceMaker: TypeAlias = Callable[[T, Sequence[str]], Sequence[Any]] class CompositeInfo(TypeInfo): @@ -84,54 +89,52 @@ WHERE t.oid = {regtype} ).format(regtype=cls._to_regtype(conn)) -class _SequenceDumper(RecursiveDumper): - def _dump_sequence( - self, obj: Sequence[Any], start: bytes, end: bytes, sep: bytes - ) -> bytes: - if not obj: - return start + end - - parts: list[abc.Buffer] = [start] +class TupleDumper(RecursiveDumper): + # Should be this, but it doesn't work + # oid = _oids.RECORD_OID - for item in obj: - if item is None: - parts.append(sep) - continue + def dump(self, obj: tuple[Any, ...]) -> Buffer | None: + return _dump_text_sequence(obj, self._tx) - dumper = self._tx.get_dumper(item, PyFormat.from_pq(self.format)) - if (ad := dumper.dump(item)) is None: - ad = b"" - elif not ad: - ad = b'""' - elif self._re_needs_quotes.search(ad): - ad = b'"' + self._re_esc.sub(rb"\1\1", ad) + b'"' - parts.append(ad) - parts.append(sep) +class _SequenceDumper(RecursiveDumper, Generic[T], ABC): + """ + Base class for text dumpers taking an object and dumping it as a composite. - parts[-1] = end + Abstract class: subclasses must specify the names of the fields making the + composite to return and the `make_sequence()` static method to convert the + object to dump to a sequence of values. + """ - return b"".join(parts) + # Subclasses must set this info + field_names: tuple[str, ...] + field_types: tuple[int, ...] - _re_needs_quotes = re.compile(rb'[",\\\s()]') - _re_esc = re.compile(rb"([\\\"])") + def dump(self, obj: T) -> bytes: + seq = type(self).make_sequence(obj, self.field_names) + return _dump_text_sequence(seq, self._tx) + @staticmethod + @abstractmethod + def make_sequence(obj: T, names: Sequence[str]) -> Sequence[Any]: ... -class TupleDumper(_SequenceDumper): - # Should be this, but it doesn't work - # oid = _oids.RECORD_OID - def dump(self, obj: tuple[Any, ...]) -> Buffer | None: - return self._dump_sequence(obj, b"(", b")", b",") +class _SequenceBinaryDumper(Dumper, Generic[T], ABC): + """ + Base class for binary dumpers taking an object and dumping it as a composite. + Abstract class: subclasses must specify the names and types of the fields + making the target composite and the `make_sequence()` static method to + convert the object to dump to a sequence of values. + """ -class _SequenceBinaryDumper(Dumper): format = pq.Format.BINARY # Subclasses must set this info + field_names: tuple[str, ...] field_types: tuple[int, ...] - def __init__(self, cls: type, context: abc.AdaptContext | None = None): + def __init__(self, cls: type[T], context: abc.AdaptContext | None = None): super().__init__(cls, context) # Note: this class is not a RecursiveDumper because it would use the @@ -144,19 +147,13 @@ class _SequenceBinaryDumper(Dumper): nfields = len(self.field_types) self._formats = (PyFormat.from_pq(self.format),) * nfields - def dump(self, obj: tuple[Any, ...]) -> Buffer | None: - out = bytearray(pack_len(len(obj))) - adapted = self._tx.dump_sequence(obj, self._formats) - for i in range(len(obj)): - b = adapted[i] - oid = self.field_types[i] - if b is not None: - out += _pack_oidlen(oid, len(b)) - out += b - else: - out += _pack_oidlen(oid, -1) + def dump(self, obj: T) -> Buffer | None: + seq = type(self).make_sequence(obj, self.field_names) + return _dump_binary_sequence(seq, self.field_types, self._formats, self._tx) - return out + @staticmethod + @abstractmethod + def make_sequence(obj: T, names: Sequence[str]) -> Sequence[Any]: ... class RecordLoader(RecursiveLoader): @@ -282,7 +279,8 @@ def register_composite( info: CompositeInfo, context: abc.AdaptContext | None = None, factory: Callable[..., T] | None = None, - make_instance: Callable[[Sequence[Any], Sequence[str]], T] | None = None, + make_instance: InstanceMaker[T] | None = None, + make_sequence: SequenceMaker[T] | None = None, ) -> None: """Register the adapters to load and dump a composite type. @@ -293,10 +291,13 @@ def register_composite( :type context: `~psycopg.abc.AdaptContext` | `!None` :param factory: Callable to convert the sequence of attributes read from the composite into a Python object. - :type factory: `!Callable[[Any, ...], T]` | `!None` + :type factory: `!Callable[..., T]` | `!None` :param make_instance: optional function taking values and names as input and returning the new type. :type make_instance: `!Callable[[Sequence[Any], Sequence[str]], T]` | `!None` + :param make_sequence: optional function taking an instance and names as + input and returning the fields to dump. + :type make_sequence: `!Callable[[T, Sequence[str]], Sequence[Any]]` | `!None` .. note:: @@ -341,16 +342,53 @@ def register_composite( # If the factory is a type, create and register dumpers for it if isinstance(factory, type): + + # Optimistically assume that the factory type is a sequence. + # If it is not, it will create a non-functioning dumper, but we don't + # risk backward incompatibility. + if not make_sequence: + + if not issubclass(factory, Sequence): + logger.warning( + "the type %r is not a sequence: dumping these objects to the" + " database will fail. Please specify a `make_sequence`" + " argument in the `register_composite()` call", + factory.__name__, + ) + + def make_sequence(obj: T, name: Sequence[str]) -> Sequence[Any]: + raise TypeError( + f"{type(obj).__name__!r} objects cannot be dumped without" + " specifying 'make_sequence' in 'register_composite()'" + ) + + else: + + def make_sequence(obj: T, name: Sequence[str]) -> Sequence[Any]: + return obj # type: ignore[return-value] # it's a sequence + + type_name = factory.__name__ dumper: type[Dumper] - dumper = _make_binary_dumper(info.name, info.oid, field_types) + dumper = _make_binary_dumper( + type_name, info.oid, field_names, field_types, make_sequence + ) adapters.register_dumper(factory, dumper) # Default to the text dumper because it is more flexible - dumper = _make_dumper(info.name, info.oid) + dumper = _make_dumper( + type_name, info.oid, field_names, field_types, make_sequence + ) adapters.register_dumper(factory, dumper) info.python_type = factory + else: + if make_sequence: + raise TypeError( + "the factory {factory.__name__!r} is not a type: you cannot" + " create a dumper by specifying `make_sequence`." + ) + def register_default_adapters(context: abc.AdaptContext) -> None: adapters = context.adapters @@ -365,6 +403,57 @@ def _nt_from_info(info: CompositeInfo) -> type[NamedTuple]: return _make_nt(name, fields) +def _dump_text_sequence(seq: Sequence[Any], tx: abc.Transformer) -> bytes: + if not seq: + return b"()" + + parts: list[abc.Buffer] = [b"("] + + for item in seq: + if item is None: + parts.append(b",") + continue + + dumper = tx.get_dumper(item, PyFormat.TEXT) + if (ad := dumper.dump(item)) is None: + ad = b"" + elif not ad: + ad = b'""' + elif _re_needs_quotes.search(ad): + ad = b'"' + _re_esc.sub(rb"\1\1", ad) + b'"' + + parts.append(ad) + parts.append(b",") + + parts[-1] = b")" + + return b"".join(parts) + + +_re_needs_quotes = re.compile(rb'[",\\\s()]') +_re_esc = re.compile(rb"([\\\"])") + + +def _dump_binary_sequence( + seq: Sequence[Any], + types: Sequence[int], + formats: Sequence[PyFormat], + tx: abc.Transformer, +) -> bytearray: + out = bytearray(pack_len(len(seq))) + adapted = tx.dump_sequence(seq, formats) + for i in range(len(seq)): + b = adapted[i] + oid = types[i] + if b is not None: + out += _pack_oidlen(oid, len(b)) + out += b + else: + out += _pack_oidlen(oid, -1) + + return out + + def _parse_text_record(data: abc.Buffer) -> list[bytes | None]: """ Split a non-empty representation of a composite type into components. @@ -438,7 +527,7 @@ def _make_loader( name: str, field_names: tuple[str, ...], field_types: tuple[int, ...], - make_instance: Callable[[Sequence[Any], Sequence[str]], T], + make_instance: InstanceMaker[T], ) -> type[_CompositeLoader[T]]: doc = f"Text loader for the '{name}' composite." d = { @@ -455,7 +544,7 @@ def _make_binary_loader( name: str, field_names: tuple[str, ...], field_types: tuple[int, ...], - make_instance: Callable[[Sequence[Any], Sequence[str]], T] | None, + make_instance: InstanceMaker[T], ) -> type[_CompositeBinaryLoader[T]]: doc = f"Binary loader for the '{name}' composite." d = { @@ -468,13 +557,38 @@ def _make_binary_loader( @cache -def _make_dumper(name: str, oid: int) -> type[TupleDumper]: - return type(f"{name.title()}Dumper", (TupleDumper,), {"oid": oid}) +def _make_dumper( + name: str, + oid: int, + field_names: tuple[str, ...], + field_types: tuple[int, ...], + make_sequence: SequenceMaker[T], +) -> type[_SequenceDumper[T]]: + doc = f"Text dumper for the '{name}' composite." + d = { + "__doc__": doc, + "oid": oid, + "field_names": field_names, + "field_types": field_types, + "make_sequence": make_sequence, + } + return type(f"{name}Dumper", (_SequenceDumper,), d) @cache def _make_binary_dumper( - name: str, oid: int, field_types: tuple[int, ...] -) -> type[_SequenceBinaryDumper]: - d = {"oid": oid, "field_types": field_types} - return type(f"{name.title()}BinaryDumper", (_SequenceBinaryDumper,), d) + name: str, + oid: int, + field_names: tuple[str, ...], + field_types: tuple[int, ...], + make_sequence: SequenceMaker[T], +) -> type[_SequenceBinaryDumper[T]]: + doc = f"Text dumper for the '{name}' composite." + d = { + "__doc__": doc, + "oid": oid, + "field_names": field_names, + "field_types": field_types, + "make_sequence": make_sequence, + } + return type(f"{name}BinaryDumper", (_SequenceBinaryDumper,), d) diff --git a/tests/types/test_composite.py b/tests/types/test_composite.py index f02da5416..de808b840 100644 --- a/tests/types/test_composite.py +++ b/tests/types/test_composite.py @@ -1,10 +1,12 @@ +import logging + import pytest from psycopg import postgres, pq, sql from psycopg.adapt import PyFormat from psycopg.postgres import types as builtins from psycopg.types.range import Range -from psycopg.types.composite import CompositeInfo, TupleDumper, register_composite +from psycopg.types.composite import CompositeInfo, register_composite from ..utils import eur from ..fix_crdb import crdb_skip_message, is_crdb @@ -344,14 +346,15 @@ def test_load_composite_factory(conn, testcomp, fmt_out): assert isinstance(res[0].baz, float) +class MyKeywordThing: + def __init__(self, *, foo, bar, baz): + self.foo, self.bar, self.baz = foo, bar, baz + + @pytest.mark.parametrize("fmt_out", pq.Format) def test_load_keyword_composite_factory(conn, testcomp, fmt_out): info = CompositeInfo.fetch(conn, "testcomp") - class MyKeywordThing: - def __init__(self, *, foo, bar, baz): - self.foo, self.bar, self.baz = foo, bar, baz - def make_instance(values, names): return MyKeywordThing(**dict(zip(names, values))) @@ -395,29 +398,20 @@ def test_register_scope(conn, testcomp): assert oid in conn.adapters._loaders[fmt] -def test_type_dumper_registered(conn, testcomp): - info = CompositeInfo.fetch(conn, "testcomp") - register_composite(info, conn) - assert issubclass(info.python_type, tuple) - assert info.python_type.__name__ == "testcomp" - d = conn.adapters.get_dumper(info.python_type, "s") - assert issubclass(d, TupleDumper) - assert d is not TupleDumper - - tc = info.python_type("foo", 42, 3.14) - cur = conn.execute("select pg_typeof(%s)", [tc]) - assert cur.fetchone()[0] == "testcomp" - - -def test_type_dumper_registered_binary(conn, testcomp): +@pytest.mark.parametrize("fmt_in", PyFormat) +def test_type_dumper_registered(conn, testcomp, fmt_in): info = CompositeInfo.fetch(conn, "testcomp") register_composite(info, conn) assert issubclass(info.python_type, tuple) assert info.python_type.__name__ == "testcomp" + assert conn.adapters.get_dumper(info.python_type, "s") tc = info.python_type("foo", 42, 3.14) - cur = conn.execute("select pg_typeof(%b)", [tc]) - assert cur.fetchone()[0] == "testcomp" + cur = conn.execute( + f"select pg_typeof(%(obj){fmt_in.value}), (%(obj){fmt_in.value}).bar", + {"obj": tc}, + ) + assert cur.fetchone() == ("testcomp", 42) def test_callable_dumper_not_registered(conn, testcomp): @@ -434,6 +428,43 @@ def test_callable_dumper_not_registered(conn, testcomp): assert cur.fetchone()[0] == ("foo", 42, 3.14, 3.14) +@pytest.mark.parametrize("fmt_in", PyFormat) +def test_dump_no_sequence(conn, testcomp, fmt_in, caplog): + caplog.set_level(logging.WARNING, logger="psycopg") + + def make_sequence(obj, names): + return [getattr(obj, attr) for attr in names] + + info = CompositeInfo.fetch(conn, "testcomp") + register_composite(info, conn, factory=MyKeywordThing, make_sequence=make_sequence) + assert info.python_type is MyKeywordThing + assert not caplog.records + + obj = MyKeywordThing(foo="foo", bar=42, baz=3.14) + cur = conn.execute( + f"select pg_typeof(%(obj){fmt_in.value}), (%(obj){fmt_in.value}).bar", + {"obj": obj}, + ) + assert cur.fetchone() == ("testcomp", 42) + + +@pytest.mark.parametrize("fmt_in", PyFormat) +def test_dump_no_sequence_failing(conn, testcomp, fmt_in, caplog): + caplog.set_level(logging.WARNING, logger="psycopg") + + info = CompositeInfo.fetch(conn, "testcomp") + register_composite(info, conn, factory=MyKeywordThing) + assert info.python_type is MyKeywordThing + assert caplog.records + assert "'MyKeywordThing' is not a sequence" in caplog.records[0].message + + obj = MyKeywordThing(foo="foo", bar=42, baz=3.14) + with pytest.raises( + TypeError, match="MyKeywordThing.*make_sequence.*register_composite" + ): + conn.execute(f"select pg_typeof(%{fmt_in.value})", [obj]) + + def test_no_info_error(conn): with pytest.raises(TypeError, match="composite"): register_composite(None, conn) # type: ignore[arg-type] -- 2.47.3