Skip to content

Commit 2a0b1bb

Browse files
elpranselprans
committed
Decode arrays as lists, handle arrays of composite types
Arrays are now decoded as lists. On the encoding side, any non-trivial container type is still accepted, but tuples are treated as composite types. This is consistent with PL/Python handling of arrays [1]. Github-Issue: #33 [1] https://www.postgresql.org/docs/devel/static/plpython-data.html#PLPYTHON-ARRAYS [2] postgres/postgres@cfd9c87
1 parent 192b6fe commit 2a0b1bb

File tree

3 files changed

+76
-52
lines changed

3 files changed

+76
-52
lines changed

‎asyncpg/protocol/codecs/array.pyx‎

Lines changed: 28 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ cdef inline _is_container(object obj):
3131
return not _is_trivial_container(obj) and isinstance(obj, ContainerABC)
3232

3333

34+
cdef inline _is_sub_array(object obj):
35+
return not _is_trivial_container(obj) and isinstance(obj, ContainerABC) \
36+
and not cpython.PyTuple_Check(obj)
37+
38+
3439
cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
3540
cdef:
3641
int32_t mylen = len(obj)
@@ -45,7 +50,7 @@ cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
4550
dims[ndims[0] - 1] = mylen
4651

4752
for elem in obj:
48-
if _is_container(elem):
53+
if _is_sub_array(elem):
4954
if elemlen == -2:
5055
elemlen = len(elem)
5156
ndims[0] += 1
@@ -133,7 +138,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
133138
int32_t ndims = hton.unpack_int32(buf.read(4))
134139
int32_t flags = hton.unpack_int32(buf.read(4))
135140
uint32_t elem_oid = hton.unpack_int32(buf.read(4))
136-
tuple result
141+
list result
137142
uint32_t i
138143
int32_t elem_len
139144
int64_t elem_count = 1
@@ -142,7 +147,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
142147
Codec elem_codec
143148

144149
if ndims == 0:
145-
result = ()
150+
result = cpython.PyList_New(0)
146151
return result
147152

148153
if ndims > ARRAY_MAXDIM:
@@ -169,7 +174,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
169174

170175
if ndims == 1:
171176
# Fast path for flat arrays
172-
result = cpython.PyTuple_New(elem_count)
177+
result = cpython.PyList_New(elem_count)
173178

174179
for i in range(elem_count):
175180
elem_len = hton.unpack_int32(buf.read(4))
@@ -180,7 +185,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
180185
elem = decoder(settings, elem_buf, decoder_arg)
181186

182187
cpython.Py_INCREF(elem)
183-
cpython.PyTuple_SET_ITEM(result, i, elem)
188+
cpython.PyList_SET_ITEM(result, i, elem)
184189

185190
else:
186191
result = _nested_array_decode(settings, buf,
@@ -200,20 +205,20 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
200205
cdef:
201206
int32_t elem_len
202207
int32_t d1, d2, d3, d4, d5, d6
203-
tuple result
208+
list result
204209
object elem
205-
tuple stride1, stride2, stride3, stride4, stride5
210+
list stride1, stride2, stride3, stride4, stride5
206211

207212
# Nested array. The approach here is dumb, but fast: rely
208213
# on the dimension limit and shape data using nested loops.
209214
# Alas, Cython doesn't have preprocessor macros.
210215
#
211-
result = cpython.PyTuple_New(dims[0])
216+
result = cpython.PyList_New(dims[0])
212217

213218
for d1 in range(dims[0]):
214-
stride1 = cpython.PyTuple_New(dims[1])
219+
stride1 = cpython.PyList_New(dims[1])
215220
cpython.Py_INCREF(stride1)
216-
cpython.PyTuple_SET_ITEM(result, d1, stride1)
221+
cpython.PyList_SET_ITEM(result, d1, stride1)
217222

218223
for d2 in range(dims[1]):
219224
if ndims == 2:
@@ -226,12 +231,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
226231
decoder_arg)
227232

228233
cpython.Py_INCREF(elem)
229-
cpython.PyTuple_SET_ITEM(stride1, d2, elem)
234+
cpython.PyList_SET_ITEM(stride1, d2, elem)
230235

231236
else:
232-
stride2 = cpython.PyTuple_New(dims[2])
237+
stride2 = cpython.PyList_New(dims[2])
233238
cpython.Py_INCREF(stride2)
234-
cpython.PyTuple_SET_ITEM(stride1, d2, stride2)
239+
cpython.PyList_SET_ITEM(stride1, d2, stride2)
235240

236241
for d3 in range(dims[2]):
237242
if ndims == 3:
@@ -244,12 +249,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
244249
decoder_arg)
245250

246251
cpython.Py_INCREF(elem)
247-
cpython.PyTuple_SET_ITEM(stride2, d3, elem)
252+
cpython.PyList_SET_ITEM(stride2, d3, elem)
248253

249254
else:
250-
stride3 = cpython.PyTuple_New(dims[3])
255+
stride3 = cpython.PyList_New(dims[3])
251256
cpython.Py_INCREF(stride3)
252-
cpython.PyTuple_SET_ITEM(stride2, d3, stride3)
257+
cpython.PyList_SET_ITEM(stride2, d3, stride3)
253258

254259
for d4 in range(dims[3]):
255260
if ndims == 4:
@@ -262,12 +267,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
262267
decoder_arg)
263268

264269
cpython.Py_INCREF(elem)
265-
cpython.PyTuple_SET_ITEM(stride3, d4, elem)
270+
cpython.PyList_SET_ITEM(stride3, d4, elem)
266271

267272
else:
268-
stride4 = cpython.PyTuple_New(dims[4])
273+
stride4 = cpython.PyList_New(dims[4])
269274
cpython.Py_INCREF(stride4)
270-
cpython.PyTuple_SET_ITEM(stride3, d4, stride4)
275+
cpython.PyList_SET_ITEM(stride3, d4, stride4)
271276

272277
for d5 in range(dims[4]):
273278
if ndims == 5:
@@ -280,12 +285,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
280285
decoder_arg)
281286

282287
cpython.Py_INCREF(elem)
283-
cpython.PyTuple_SET_ITEM(stride4, d5, elem)
288+
cpython.PyList_SET_ITEM(stride4, d5, elem)
284289

285290
else:
286-
stride5 = cpython.PyTuple_New(dims[5])
291+
stride5 = cpython.PyList_New(dims[5])
287292
cpython.Py_INCREF(stride5)
288-
cpython.PyTuple_SET_ITEM(stride4, d5, stride5)
293+
cpython.PyList_SET_ITEM(stride4, d5, stride5)
289294

290295
for d6 in range(dims[5]):
291296
elem_len = hton.unpack_int32(buf.read(4))
@@ -297,7 +302,7 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
297302
decoder_arg)
298303

299304
cpython.Py_INCREF(elem)
300-
cpython.PyTuple_SET_ITEM(stride5, d6, elem)
305+
cpython.PyList_SET_ITEM(stride5, d6, elem)
301306

302307
return result
303308

‎asyncpg/protocol/codecs/base.pyx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,10 @@ cdef class DataCodecConfig:
313313
schema = ti['ns']
314314
array_element_oid = ti['elemtype']
315315
range_subtype_oid = ti['range_subtype']
316-
comp_type_attrs = ti['attrtypoids']
316+
if ti['attrtypoids']:
317+
comp_type_attrs = tuple(ti['attrtypoids'])
318+
else:
319+
comp_type_attrs = None
317320
base_type = ti['basetype']
318321

319322
if array_element_oid:

‎tests/test_codecs.py‎

Lines changed: 44 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -202,9 +202,9 @@ def _timezone(offset):
202202
'output': uuid.UUID('00000000-0000-0000-0000-000000000000')}
203203
]),
204204
('uuid[]', 'uuid[]', [
205-
(uuid.UUID('38a4ff5a-3a56-11e6-a6c2-c8f73323c6d4'),
206-
uuid.UUID('00000000-0000-0000-0000-000000000000')),
207-
()
205+
[uuid.UUID('38a4ff5a-3a56-11e6-a6c2-c8f73323c6d4'),
206+
uuid.UUID('00000000-0000-0000-0000-000000000000')],
207+
[]
208208
]),
209209
('json', 'json', [
210210
'[1, 2, 3, 4]',
@@ -215,30 +215,30 @@ def _timezone(offset):
215215
'{"a": [1, 2], "b": 0}'
216216
]),
217217
('oid[]', 'oid[]', [
218-
(1, 2, 3, 4),
219-
()
218+
[1, 2, 3, 4],
219+
[]
220220
]),
221221
('smallint[]', 'int2[]', [
222-
(1, 2, 3, 4),
223-
(1, 2, 3, 4, 5, 6, 7, 8, 9, 0),
224-
()
222+
[1, 2, 3, 4],
223+
[1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
224+
[]
225225
]),
226226
('bigint[]', 'int8[]', [
227-
(2 ** 42, -2 ** 54, 0),
228-
()
227+
[2 ** 42, -2 ** 54, 0],
228+
[]
229229
]),
230230
('int[]', 'int4[]', [
231-
(2 ** 22, -2 ** 24, 0),
232-
()
231+
[2 ** 22, -2 ** 24, 0],
232+
[]
233233
]),
234234
('time[]', 'time[]', [
235-
(datetime.time(12, 15, 20), datetime.time(0, 1, 1)),
236-
()
235+
[datetime.time(12, 15, 20), datetime.time(0, 1, 1)],
236+
[]
237237
]),
238238
('text[]', 'text[]', [
239-
('ABCDE', 'EDCBA'),
240-
(),
241-
('A' * 1024 * 1024,) * 10
239+
['ABCDE', 'EDCBA'],
240+
[],
241+
['A' * 1024 * 1024] * 10
242242
]),
243243
('float8', 'float8', [
244244
1.1,
@@ -474,15 +474,15 @@ async def test_arrays(self):
474474
cases = [
475475
(
476476
r"SELECT '[1:3][-1:0]={{1,2},{4,5},{6,7}}'::int[]",
477-
((1, 2), (4, 5), (6, 7))
477+
[[1, 2], [4, 5], [6, 7]]
478478
),
479479
(
480480
r"SELECT '{{{{{{1}}}}}}'::int[]",
481-
((((((1,),),),),),)
481+
[[[[[[1]]]]]]
482482
),
483483
(
484484
r"SELECT '{1, 2, NULL}'::int[]::anyarray",
485-
(1, 2, None)
485+
[1, 2, None]
486486
),
487487
]
488488

@@ -495,12 +495,12 @@ async def test_arrays(self):
495495
await self.con.fetchval("SELECT '{{{{{{{1}}}}}}}'::int[]")
496496

497497
cases = [
498-
(None,),
499-
(1, 2, 3, 4, 5, 6),
500-
((1, 2), (4, 5), (6, 7)),
501-
(((1,), (2,)), ((4,), (5,)), ((None,), (7,))),
502-
((((((1,),),),),),),
503-
((((((None,),),),),),)
498+
[None],
499+
[1, 2, 3, 4, 5, 6],
500+
[[1, 2], [4, 5], [6, 7]],
501+
[[[1], [2]], [[4], [5]], [[None], [7]]],
502+
[[[[[[1]]]]]],
503+
[[[[[[None]]]]]]
504504
]
505505

506506
st = await self.con.prepare(
@@ -579,11 +579,11 @@ async def test_composites(self):
579579

580580
self.assertIsNone(res['a'])
581581
self.assertEqual(res['b'], '5678')
582-
self.assertEqual(res['c'], (9, None, 11))
582+
self.assertEqual(res['c'], [9, None, 11])
583583

584584
self.assertIsNone(res[0])
585585
self.assertEqual(res[1], '5678')
586-
self.assertEqual(res[2], (9, None, 11))
586+
self.assertEqual(res[2], [9, None, 11])
587587

588588
at = st.get_attributes()
589589
self.assertEqual(len(at), 1)
@@ -860,3 +860,19 @@ def hstore_encoder(obj):
860860
await self.con.execute('''
861861
DROP EXTENSION hstore
862862
''')
863+
864+
async def test_composites_in_arrays(self):
865+
await self.con.execute('''
866+
CREATE TYPE t AS (a text, b int);
867+
CREATE TABLE tab (d t[]);
868+
''')
869+
870+
await self.con.execute(
871+
'INSERT INTO tab (d) VALUES ($1)',
872+
[('a', 1)])
873+
874+
r = await self.con.fetchval('''
875+
SELECT d FROM tab
876+
''')
877+
878+
self.assertEqual(r, [('a', 1)])

0 commit comments

Comments
 (0)