Allow any non-trivial Container to be passed as an array argument

Author: elprans

asyncpg/protocol/codecs/array.pyx

@@ -5,6 +5,9 @@
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
 
+from collections.abc import Container as ContainerABC
+
+
 DEF ARRAY_MAXDIM = 6  # defined in postgresql/src/includes/c.h
 
 
@@ -19,8 +22,13 @@ ctypedef object (*decode_func_ex)(ConnectionSettings settings,
                                   const void *arg)
 
 
-cdef inline bint _is_array(object obj):
-    return cpython.PyTuple_Check(obj) or cpython.PyList_Check(obj)
+cdef inline bint _is_trivial_container(object obj):
+    return cpython.PyUnicode_Check(obj) or cpython.PyBytes_Check(obj) or \
+            PyByteArray_Check(obj) or PyMemoryView_Check(obj)
+
+
+cdef inline _is_container(object obj):
+    return not _is_trivial_container(obj) and isinstance(obj, ContainerABC)
 
 
 cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
@@ -37,7 +45,7 @@ cdef _get_array_shape(object obj, int32_t *dims, int32_t *ndims):
     dims[ndims[0] - 1] = mylen
 
     for elem in obj:
-        if _is_array(elem):
+        if _is_container(elem):
             if elemlen == -2:
                 elemlen = len(elem)
                 ndims[0] += 1
@@ -80,9 +88,10 @@ cdef inline array_encode(ConnectionSettings settings, WriteBuffer buf,
         int32_t ndims = 1
         int32_t i
 
-    if not _is_array(obj):
+    if not _is_container(obj):
         raise TypeError(
-            'list or tuple expected (got type {})'.format(type(obj)))
+            'a non-trivial iterable expected (got type {!r})'.format(
+                type(obj).__name__))
 
     _get_array_shape(obj, dims, &ndims)
 

asyncpg/protocol/protocol.pyx

@@ -23,7 +23,8 @@ from asyncpg.protocol cimport record
 from asyncpg.protocol.python cimport (
                      PyMem_Malloc, PyMem_Realloc, PyMem_Calloc, PyMem_Free,
                      PyMemoryView_GET_BUFFER, PyMemoryView_Check,
-                     PyUnicode_AsUTF8AndSize, PyByteArray_AsString)
+                     PyUnicode_AsUTF8AndSize, PyByteArray_AsString,
+                     PyByteArray_Check)
 
 from cpython cimport PyBuffer_FillInfo, PyBytes_AsString
 

asyncpg/protocol/python.pxd

@@ -13,7 +13,10 @@ cdef extern from "Python.h":
     void* PyMem_Calloc(size_t nelem, size_t elsize)  # Python >= 3.5!
     void PyMem_Free(void *p)
 
+    int PyByteArray_Check(object)
+
     int PyMemoryView_Check(object)
     Py_buffer *PyMemoryView_GET_BUFFER(object)
+
     char* PyUnicode_AsUTF8AndSize(object unicode, ssize_t *size) except NULL
     char* PyByteArray_AsString(object)

tests/test_codecs.py

@@ -542,7 +542,7 @@ async def test_arrays(self):
                 "SELECT $1::int[]",
                 [[1], ['t'], [2]])
 
-        with self.assertRaisesRegex(TypeError, 'list or tuple expected'):
+        with self.assertRaisesRegex(TypeError, 'non-trivial iterable expected'):
             await self.con.fetchval(
                 "SELECT $1::int[]",
                 1)