Skip to content

Commit eaa2fa1

Browse files
committed
Allow mappings as composite type input
This allows asyncpg to take an arbitrary mapping as input for composite types, as long as the keys match composite type fields. This allows removing some tedium when handling complex composite types. Fixes: #349.
1 parent df7830f commit eaa2fa1

File tree

5 files changed

+60
-6
lines changed

5 files changed

+60
-6
lines changed

asyncpg/protocol/codecs/array.pyx

+5-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@
55
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
66

77

8-
from collections.abc import Iterable as IterableABC, Sized as SizedABC
8+
from collections.abc import (Iterable as IterableABC,
9+
Mapping as MappingABC,
10+
Sized as SizedABC)
911

1012
from asyncpg import exceptions
1113

@@ -36,7 +38,8 @@ cdef inline _is_array_iterable(object obj):
3638
return (
3739
isinstance(obj, IterableABC) and
3840
isinstance(obj, SizedABC) and
39-
not _is_trivial_container(obj)
41+
not _is_trivial_container(obj) and
42+
not isinstance(obj, MappingABC)
4043
)
4144

4245

asyncpg/protocol/codecs/base.pxd

+1
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ cdef class Codec:
6767
# composite types
6868
tuple element_type_oids
6969
object element_names
70+
object record_desc
7071
list element_codecs
7172

7273
# Pointers to actual encoder/decoder functions for this codec

asyncpg/protocol/codecs/base.pyx

+31-3
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
66

77

8+
from collections.abc import Mapping as MappingABC
9+
810
from asyncpg import exceptions
911

1012

@@ -43,12 +45,13 @@ cdef class Codec:
4345
self.element_type_oids = element_type_oids
4446
self.element_codecs = element_codecs
4547
self.element_delimiter = element_delimiter
48+
self.element_names = element_names
4649

4750
if element_names is not None:
48-
self.element_names = record.ApgRecordDesc_New(
51+
self.record_desc = record.ApgRecordDesc_New(
4952
element_names, tuple(element_names))
5053
else:
51-
self.element_names = None
54+
self.record_desc = None
5255

5356
if type == CODEC_C:
5457
self.encoder = <codec_encode_func>&self.encode_scalar
@@ -125,6 +128,31 @@ cdef class Codec:
125128
int i
126129
list elem_codecs = self.element_codecs
127130
ssize_t count
131+
ssize_t composite_size
132+
tuple rec
133+
134+
if isinstance(obj, MappingABC):
135+
# Input is dict-like, form a tuple
136+
composite_size = len(self.element_type_oids)
137+
rec = cpython.PyTuple_New(composite_size)
138+
139+
for i in range(composite_size):
140+
cpython.Py_INCREF(None)
141+
cpython.PyTuple_SET_ITEM(rec, i, None)
142+
143+
for field in obj:
144+
try:
145+
i = self.element_names[field]
146+
except KeyError:
147+
raise ValueError(
148+
'{!r} is not a valid element of composite '
149+
'type {}'.format(field, self.name)) from None
150+
151+
item = obj[field]
152+
cpython.Py_INCREF(item)
153+
cpython.PyTuple_SET_ITEM(rec, i, item)
154+
155+
obj = rec
128156

129157
count = len(obj)
130158
if count > _MAXINT32:
@@ -204,7 +232,7 @@ cdef class Codec:
204232
schema=self.schema,
205233
data_type=self.name,
206234
)
207-
result = record.ApgRecord_New(self.element_names, elem_count)
235+
result = record.ApgRecord_New(self.record_desc, elem_count)
208236
for i in range(elem_count):
209237
elem_typ = self.element_type_oids[i]
210238
received_elem_typ = <uint32_t>hton.unpack_int32(buf.read(4))

docs/usage.rst

+2-1
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,8 @@ The table below shows the correspondence between PostgreSQL and Python types.
7676
| ``anyrange`` | :class:`asyncpg.Range <asyncpg.types.Range>` |
7777
+----------------------+-----------------------------------------------------+
7878
| ``record`` | :class:`asyncpg.Record`, |
79-
| | :class:`tuple <python:tuple>` |
79+
| | :class:`tuple <python:tuple>`, |
80+
| | :class:`Mapping <python:collections.abc.Mapping>` |
8081
+----------------------+-----------------------------------------------------+
8182
| ``bit``, ``varbit`` | :class:`asyncpg.BitString <asyncpg.types.BitString>`|
8283
+----------------------+-----------------------------------------------------+

tests/test_codecs.py

+21
Original file line numberDiff line numberDiff line change
@@ -851,6 +851,27 @@ async def test_composites(self):
851851
SELECT $1::test_composite
852852
''', res)
853853

854+
# composite input as a mapping
855+
res = await self.con.fetchval('''
856+
SELECT $1::test_composite
857+
''', {'b': 'foo', 'a': 1, 'c': [1, 2, 3]})
858+
859+
self.assertEqual(res, (1, 'foo', [1, 2, 3]))
860+
861+
# Test None padding
862+
res = await self.con.fetchval('''
863+
SELECT $1::test_composite
864+
''', {'a': 1})
865+
866+
self.assertEqual(res, (1, None, None))
867+
868+
with self.assertRaisesRegex(
869+
asyncpg.DataError,
870+
"'bad' is not a valid element"):
871+
await self.con.fetchval(
872+
"SELECT $1::test_composite",
873+
{'bad': 'foo'})
874+
854875
finally:
855876
await self.con.execute('DROP TYPE test_composite')
856877

0 commit comments

Comments
 (0)