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

Author: elprans

asyncpg/protocol/codecs/array.pyx

@@ -31,6 +31,11 @@ cdef inline _is_container(object obj):
     return not _is_trivial_container(obj) and isinstance(obj, ContainerABC)
 
 
+cdef inline _is_sub_array(object obj):
+    return not _is_trivial_container(obj) and isinstance(obj, ContainerABC) \
+            and not cpython.PyTuple_Check(obj)
+
+
 cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
     cdef:
         int32_t mylen = len(obj)
@@ -45,7 +50,7 @@ cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
     dims[ndims[0] - 1] = mylen
 
     for elem in obj:
-        if _is_container(elem):
+        if _is_sub_array(elem):
             if elemlen == -2:
                 elemlen = len(elem)
                 ndims[0] += 1
@@ -133,7 +138,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
         int32_t ndims = hton.unpack_int32(buf.read(4))
         int32_t flags = hton.unpack_int32(buf.read(4))
         uint32_t elem_oid = hton.unpack_int32(buf.read(4))
-        tuple result
+        list result
         uint32_t i
         int32_t elem_len
         int64_t elem_count = 1
@@ -142,7 +147,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
         Codec elem_codec
 
     if ndims == 0:
-        result = ()
+        result = cpython.PyList_New(0)
         return result
 
     if ndims > ARRAY_MAXDIM:
@@ -169,7 +174,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
 
     if ndims == 1:
         # Fast path for flat arrays
-        result = cpython.PyTuple_New(elem_count)
+        result = cpython.PyList_New(elem_count)
 
         for i in range(elem_count):
             elem_len = hton.unpack_int32(buf.read(4))
@@ -180,7 +185,7 @@ cdef inline array_decode(ConnectionSettings settings, FastReadBuffer buf,
                 elem = decoder(settings, elem_buf, decoder_arg)
 
             cpython.Py_INCREF(elem)
-            cpython.PyTuple_SET_ITEM(result, i, elem)
+            cpython.PyList_SET_ITEM(result, i, elem)
 
     else:
         result = _nested_array_decode(settings, buf,
@@ -200,20 +205,20 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
     cdef:
         int32_t elem_len
         int32_t d1, d2, d3, d4, d5, d6
-        tuple result
+        list result
         object elem
-        tuple stride1, stride2, stride3, stride4, stride5
+        list stride1, stride2, stride3, stride4, stride5
 
     # Nested array.  The approach here is dumb, but fast: rely
     # on the dimension limit and shape data using nested loops.
     # Alas, Cython doesn't have preprocessor macros.
     #
-    result = cpython.PyTuple_New(dims[0])
+    result = cpython.PyList_New(dims[0])
 
     for d1 in range(dims[0]):
-        stride1 = cpython.PyTuple_New(dims[1])
+        stride1 = cpython.PyList_New(dims[1])
         cpython.Py_INCREF(stride1)
-        cpython.PyTuple_SET_ITEM(result, d1, stride1)
+        cpython.PyList_SET_ITEM(result, d1, stride1)
 
         for d2 in range(dims[1]):
             if ndims == 2:
@@ -226,12 +231,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
                                    decoder_arg)
 
                 cpython.Py_INCREF(elem)
-                cpython.PyTuple_SET_ITEM(stride1, d2, elem)
+                cpython.PyList_SET_ITEM(stride1, d2, elem)
 
             else:
-                stride2 = cpython.PyTuple_New(dims[2])
+                stride2 = cpython.PyList_New(dims[2])
                 cpython.Py_INCREF(stride2)
-                cpython.PyTuple_SET_ITEM(stride1, d2, stride2)
+                cpython.PyList_SET_ITEM(stride1, d2, stride2)
 
                 for d3 in range(dims[2]):
                     if ndims == 3:
@@ -244,12 +249,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
                                            decoder_arg)
 
                         cpython.Py_INCREF(elem)
-                        cpython.PyTuple_SET_ITEM(stride2, d3, elem)
+                        cpython.PyList_SET_ITEM(stride2, d3, elem)
 
                     else:
-                        stride3 = cpython.PyTuple_New(dims[3])
+                        stride3 = cpython.PyList_New(dims[3])
                         cpython.Py_INCREF(stride3)
-                        cpython.PyTuple_SET_ITEM(stride2, d3, stride3)
+                        cpython.PyList_SET_ITEM(stride2, d3, stride3)
 
                         for d4 in range(dims[3]):
                             if ndims == 4:
@@ -262,12 +267,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
                                                    decoder_arg)
 
                                 cpython.Py_INCREF(elem)
-                                cpython.PyTuple_SET_ITEM(stride3, d4, elem)
+                                cpython.PyList_SET_ITEM(stride3, d4, elem)
 
                             else:
-                                stride4 = cpython.PyTuple_New(dims[4])
+                                stride4 = cpython.PyList_New(dims[4])
                                 cpython.Py_INCREF(stride4)
-                                cpython.PyTuple_SET_ITEM(stride3, d4, stride4)
+                                cpython.PyList_SET_ITEM(stride3, d4, stride4)
 
                                 for d5 in range(dims[4]):
                                     if ndims == 5:
@@ -280,12 +285,12 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
                                                            decoder_arg)
 
                                         cpython.Py_INCREF(elem)
-                                        cpython.PyTuple_SET_ITEM(stride4, d5, elem)
+                                        cpython.PyList_SET_ITEM(stride4, d5, elem)
 
                                     else:
-                                        stride5 = cpython.PyTuple_New(dims[5])
+                                        stride5 = cpython.PyList_New(dims[5])
                                         cpython.Py_INCREF(stride5)
-                                        cpython.PyTuple_SET_ITEM(stride4, d5, stride5)
+                                        cpython.PyList_SET_ITEM(stride4, d5, stride5)
 
                                         for d6 in range(dims[5]):
                                             elem_len = hton.unpack_int32(buf.read(4))
@@ -297,7 +302,7 @@ cdef inline _nested_array_decode(ConnectionSettings settings,
                                                                decoder_arg)
 
                                             cpython.Py_INCREF(elem)
-                                            cpython.PyTuple_SET_ITEM(stride5, d6, elem)
+                                            cpython.PyList_SET_ITEM(stride5, d6, elem)
 
     return result
 

asyncpg/protocol/codecs/base.pyx

@@ -313,7 +313,10 @@ cdef class DataCodecConfig:
             schema = ti['ns']
             array_element_oid = ti['elemtype']
             range_subtype_oid = ti['range_subtype']
-            comp_type_attrs = ti['attrtypoids']
+            if ti['attrtypoids']:
+                comp_type_attrs = tuple(ti['attrtypoids'])
+            else:
+                comp_type_attrs = None
             base_type = ti['basetype']
 
             if array_element_oid:

tests/test_codecs.py

@@ -202,9 +202,9 @@ def _timezone(offset):
          'output': uuid.UUID('00000000-0000-0000-0000-000000000000')}
     ]),
     ('uuid[]', 'uuid[]', [
-        (uuid.UUID('38a4ff5a-3a56-11e6-a6c2-c8f73323c6d4'),
-         uuid.UUID('00000000-0000-0000-0000-000000000000')),
-        ()
+        [uuid.UUID('38a4ff5a-3a56-11e6-a6c2-c8f73323c6d4'),
+         uuid.UUID('00000000-0000-0000-0000-000000000000')],
+        []
     ]),
     ('json', 'json', [
         '[1, 2, 3, 4]',
@@ -215,30 +215,30 @@ def _timezone(offset):
         '{"a": [1, 2], "b": 0}'
     ]),
     ('oid[]', 'oid[]', [
-        (1, 2, 3, 4),
-        ()
+        [1, 2, 3, 4],
+        []
     ]),
     ('smallint[]', 'int2[]', [
-        (1, 2, 3, 4),
-        (1, 2, 3, 4, 5, 6, 7, 8, 9, 0),
-        ()
+        [1, 2, 3, 4],
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
+        []
     ]),
     ('bigint[]', 'int8[]', [
-        (2 ** 42, -2 ** 54, 0),
-        ()
+        [2 ** 42, -2 ** 54, 0],
+        []
     ]),
     ('int[]', 'int4[]', [
-        (2 ** 22, -2 ** 24, 0),
-        ()
+        [2 ** 22, -2 ** 24, 0],
+        []
     ]),
     ('time[]', 'time[]', [
-        (datetime.time(12, 15, 20), datetime.time(0, 1, 1)),
-        ()
+        [datetime.time(12, 15, 20), datetime.time(0, 1, 1)],
+        []
     ]),
     ('text[]', 'text[]', [
-        ('ABCDE', 'EDCBA'),
-        (),
-        ('A' * 1024 * 1024,) * 10
+        ['ABCDE', 'EDCBA'],
+        [],
+        ['A' * 1024 * 1024] * 10
     ]),
     ('float8', 'float8', [
         1.1,
@@ -474,15 +474,15 @@ async def test_arrays(self):
         cases = [
             (
                 r"SELECT '[1:3][-1:0]={{1,2},{4,5},{6,7}}'::int[]",
-                ((1, 2), (4, 5), (6, 7))
+                [[1, 2], [4, 5], [6, 7]]
             ),
             (
                 r"SELECT '{{{{{{1}}}}}}'::int[]",
-                ((((((1,),),),),),)
+                [[[[[[1]]]]]]
             ),
             (
                 r"SELECT '{1, 2, NULL}'::int[]::anyarray",
-                (1, 2, None)
+                [1, 2, None]
             ),
         ]
 
@@ -495,12 +495,12 @@ async def test_arrays(self):
             await self.con.fetchval("SELECT '{{{{{{{1}}}}}}}'::int[]")
 
         cases = [
-            (None,),
-            (1, 2, 3, 4, 5, 6),
-            ((1, 2), (4, 5), (6, 7)),
-            (((1,), (2,)), ((4,), (5,)), ((None,), (7,))),
-            ((((((1,),),),),),),
-            ((((((None,),),),),),)
+            [None],
+            [1, 2, 3, 4, 5, 6],
+            [[1, 2], [4, 5], [6, 7]],
+            [[[1], [2]], [[4], [5]], [[None], [7]]],
+            [[[[[[1]]]]]],
+            [[[[[[None]]]]]]
         ]
 
         st = await self.con.prepare(
@@ -579,11 +579,11 @@ async def test_composites(self):
 
             self.assertIsNone(res['a'])
             self.assertEqual(res['b'], '5678')
-            self.assertEqual(res['c'], (9, None, 11))
+            self.assertEqual(res['c'], [9, None, 11])
 
             self.assertIsNone(res[0])
             self.assertEqual(res[1], '5678')
-            self.assertEqual(res[2], (9, None, 11))
+            self.assertEqual(res[2], [9, None, 11])
 
             at = st.get_attributes()
             self.assertEqual(len(at), 1)
@@ -860,3 +860,19 @@ def hstore_encoder(obj):
             await self.con.execute('''
                 DROP EXTENSION hstore
             ''')
+
+    async def test_composites_in_arrays(self):
+        await self.con.execute('''
+            CREATE TYPE t AS (a text, b int);
+            CREATE TABLE tab (d t[]);
+        ''')
+
+        await self.con.execute(
+            'INSERT INTO tab (d) VALUES ($1)',
+            [('a', 1)])
+
+        r = await self.con.fetchval('''
+            SELECT d FROM tab
+        ''')
+
+        self.assertEqual(r, [('a', 1)])