Merge branch 'array-typecast' into devel

NEWS

@@ -1,6 +1,9 @@
 What's new in psycopg 2.4.3
 ---------------------------
 
+  - Added 'new_array_type()' function for easy creation of array
+    typecasters.
+  - Added support for arrays of hstores and composite types (ticket #66).
   - Fixed segfault in case of transaction started with connection lost
     (and possibly other events).
   - Rollback connections in transaction or in error before putting them
@@ -11,6 +14,7 @@ What's new in psycopg 2.4.3
   - Fixed interaction between RealDictCursor and named cursors
     (ticket #67).
   - Dropped limit on the columns length in COPY operations (ticket #68).
+  - Fixed typecasting of arrays containing consecutive backslashes.
   - 'errorcodes' map updated to PostgreSQL 9.1.
 
 

doc/src/advanced.rst

@@ -216,6 +216,9 @@ read:
     >>> print type(point), point.x, point.y
     <class 'Point'> 10.2 20.3
 
+A typecaster created by `!new_type()` can be also used with
+`~psycopg2.extensions.new_array_type()` to create a typecaster converting a
+PostgreSQL array into a Python list.
 
 
 .. index::

doc/src/extensions.rst

@@ -290,7 +290,7 @@ details.
 .. function:: new_type(oids, name, adapter)
 
     Create a new type caster to convert from a PostgreSQL type to a Python
-    object.  The created object must be registered using
+    object.  The object created must be registered using
     `register_type()` to be used.
 
     :param oids: tuple of OIDs of the PostgreSQL type to convert.
@@ -309,6 +309,23 @@ details.
     See :ref:`type-casting-from-sql-to-python` for an usage example.
 
 
+.. function:: new_array_type(oids, name, base_caster)
+
+    Create a new type caster to convert from a PostgreSQL array type to a list
+    of Python object.  The object created must be registered using
+    `register_type()` to be used.
+
+    :param oids: tuple of OIDs of the PostgreSQL type to convert. It should
+        probably be the oid of the array type (e.g. the ``typarray`` field in
+        the ``pg_type`` table.
+    :param name: the name of the new type adapter.
+    :param base_caster: a Psycopg typecaster, e.g. created using the
+        `new_type()` function. The caster should be able to parse a single
+        item of the desired type.
+
+    .. versionadded:: 2.4.3
+
+
 .. function:: register_type(obj [, scope])
 
     Register a type caster created using `new_type()`.

lib/extensions.py

@@ -57,7 +57,7 @@ except ImportError:
     pass
 
 from psycopg2._psycopg import adapt, adapters, encodings, connection, cursor, lobject, Xid
-from psycopg2._psycopg import string_types, binary_types, new_type, register_type
+from psycopg2._psycopg import string_types, binary_types, new_type, new_array_type, register_type
 from psycopg2._psycopg import ISQLQuote, Notify
 
 from psycopg2._psycopg import QueryCanceledError, TransactionRollbackError

lib/extras.py

@@ -699,7 +699,8 @@ WHERE typname = 'hstore';
 
         return tuple(rv0), tuple(rv1)
 
-def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
+def register_hstore(conn_or_curs, globally=False, unicode=False,
+        oid=None, array_oid=None):
     """Register adapter and typecaster for `!dict`\-\ |hstore| conversions.
 
     :param conn_or_curs: a connection or cursor: the typecaster will be
@@ -709,14 +710,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
         will be `!unicode` instead of `!str`. The option is not available on
         Python 3
     :param oid: the OID of the |hstore| type if known. If not, it will be
-        queried on *conn_or_curs*
+        queried on *conn_or_curs*.
+    :param array_oid: the OID of the |hstore| array type if known. If not, it
+        will be queried on *conn_or_curs*.
 
     The connection or cursor passed to the function will be used to query the
     database and look for the OID of the |hstore| type (which may be different
     across databases). If querying is not desirable (e.g. with
     :ref:`asynchronous connections <async-support>`) you may specify it in the
-    *oid* parameter (it can be found using a query such as :sql:`SELECT
+    *oid* parameter, which can be found using a query such as :sql:`SELECT
-    'hstore'::regtype::oid;`).
+    'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid*
+    using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`.
+
 
     Note that, when passing a dictionary from Python to the database, both
     strings and unicode keys and values are supported. Dictionaries returned
@@ -730,6 +735,10 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
         added the *oid* parameter. If not specified, the typecaster is
         installed also if |hstore| is not installed in the :sql:`public`
         schema.
+
+    .. versionchanged:: 2.4.3
+        added support for |hstore| array.
+
     """
     if oid is None:
         oid = HstoreAdapter.get_oids(conn_or_curs)
@@ -738,11 +747,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
                 "hstore type not found in the database. "
                 "please install it from your 'contrib/hstore.sql' file")
         else:
-            oid = oid[0]  # for the moment we don't have a HSTOREARRAY
+            array_oid = oid[1]
+            oid = oid[0]
 
     if isinstance(oid, int):
         oid = (oid,)
 
+    if array_oid is not None:
+        if isinstance(array_oid, int):
+            array_oid = (array_oid,)
+        else:
+            array_oid = tuple([x for x in array_oid if x])
+
     # create and register the typecaster
     if sys.version_info[0] < 3 and unicode:
         cast = HstoreAdapter.parse_unicode
@@ -753,11 +769,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
     _ext.register_type(HSTORE, not globally and conn_or_curs or None)
     _ext.register_adapter(dict, HstoreAdapter)
 
+    if array_oid:
+        HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE)
+        _ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)
+
 
 class CompositeCaster(object):
     """Helps conversion of a PostgreSQL composite type into a Python object.
 
     The class is usually created by the `register_composite()` function.
+    You may want to create and register manually instances of the class if
+    querying the database at registration time is not desirable (such as when
+    using an :ref:`asynchronous connections <async-support>`).
 
     .. attribute:: name
 
@@ -767,6 +790,10 @@ class CompositeCaster(object):
 
         The oid of the PostgreSQL type.
 
+    .. attribute:: array_oid
+
+        The oid of the PostgreSQL array type, if available.
+
     .. attribute:: type
 
         The type of the Python objects returned. If :py:func:`collections.namedtuple()`
@@ -782,14 +809,20 @@ class CompositeCaster(object):
         List of component type oids of the type to be casted.
 
     """
-    def __init__(self, name, oid, attrs):
+    def __init__(self, name, oid, attrs, array_oid=None):
         self.name = name
         self.oid = oid
+        self.array_oid = array_oid
 
         self.attnames = [ a[0] for a in attrs ]
         self.atttypes = [ a[1] for a in attrs ]
         self._create_type(name, self.attnames)
         self.typecaster = _ext.new_type((oid,), name, self.parse)
+        if array_oid:
+            self.array_typecaster = _ext.new_array_type(
+                (array_oid,), "%sARRAY" % name, self.typecaster)
+        else:
+            self.array_typecaster = None
 
     def parse(self, s, curs):
         if s is None:
@@ -861,15 +894,18 @@ class CompositeCaster(object):
             tname = name
             schema = 'public'
 
+        # column typarray not available before PG 8.3
+        typarray = conn.server_version >= 80300 and "typarray" or "NULL"
+
         # get the type oid and attributes
         curs.execute("""\
-SELECT t.oid, attname, atttypid
+SELECT t.oid, %s, attname, atttypid
 FROM pg_type t
 JOIN pg_namespace ns ON typnamespace = ns.oid
 JOIN pg_attribute a ON attrelid = typrelid
-WHERE typname = %s and nspname = %s
+WHERE typname = %%s and nspname = %%s
 ORDER BY attnum;
-""", (tname, schema))
+""" % typarray, (tname, schema))
 
         recs = curs.fetchall()
 
@@ -883,9 +919,11 @@ ORDER BY attnum;
                 "PostgreSQL type '%s' not found" % name)
 
         type_oid = recs[0][0]
-        type_attrs = [ (r[1], r[2]) for r in recs ]
+        array_oid = recs[0][1]
+        type_attrs = [ (r[2], r[3]) for r in recs ]
 
-        return CompositeCaster(tname, type_oid, type_attrs)
+        return CompositeCaster(tname, type_oid, type_attrs,
+            array_oid=array_oid)
 
 def register_composite(name, conn_or_curs, globally=False):
     """Register a typecaster to convert a composite type into a tuple.
@@ -899,10 +937,17 @@ def register_composite(name, conn_or_curs, globally=False):
         *conn_or_curs*, otherwise register it globally
     :return: the registered `CompositeCaster` instance responsible for the
         conversion
+
+    .. versionchanged:: 2.4.3
+        added support for array of composite types
+
     """
     caster = CompositeCaster._from_db(name, conn_or_curs)
     _ext.register_type(caster.typecaster, not globally and conn_or_curs or None)
 
+    if caster.array_typecaster is not None:
+        _ext.register_type(caster.array_typecaster, not globally and conn_or_curs or None)
+
     return caster
 
 

psycopg/psycopgmodule.c

@@ -257,7 +257,7 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds)
 "  * `conn_or_curs`: A connection, cursor or None"
 
 #define typecast_from_python_doc \
-"new_type(oids, name, adapter) -> new type object\n\n" \
+"new_type(oids, name, castobj) -> new type object\n\n" \
 "Create a new binding object. The object can be used with the\n" \
 "`register_type()` function to bind PostgreSQL objects to python objects.\n\n" \
 ":Parameters:\n" \
@@ -268,6 +268,15 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds)
 "    the string representation returned by PostgreSQL (`!None` if ``NULL``)\n" \
 "    and ``cur`` is the cursor from which data are read."
 
+#define typecast_array_from_python_doc \
+"new_array_type(oids, name, baseobj) -> new type object\n\n" \
+"Create a new binding object to parse an array.\n\n" \
+"The object can be used with `register_type()`.\n\n" \
+":Parameters:\n" \
+"  * `oids`: Tuple of ``oid`` of the PostgreSQL types to convert.\n" \
+"  * `name`: Name for the new type\n" \
+"  * `baseobj`: Adapter to perform type conversion of a single array item."
+
 static void
 _psyco_register_type_set(PyObject **dict, PyObject *type)
 {
@@ -758,6 +767,8 @@ static PyMethodDef psycopgMethods[] = {
      METH_VARARGS, psyco_register_type_doc},
     {"new_type", (PyCFunction)typecast_from_python,
      METH_VARARGS|METH_KEYWORDS, typecast_from_python_doc},
+    {"new_array_type", (PyCFunction)typecast_array_from_python,
+     METH_VARARGS|METH_KEYWORDS, typecast_array_from_python_doc},
 
     {"AsIs",  (PyCFunction)psyco_AsIs,
      METH_VARARGS, psyco_AsIs_doc},

psycopg/typecast.c

@@ -603,6 +603,29 @@ typecast_from_python(PyObject *self, PyObject *args, PyObject *keywds)
     return typecast_new(name, v, cast, base);
 }
 
+PyObject *
+typecast_array_from_python(PyObject *self, PyObject *args, PyObject *keywds)
+{
+    PyObject *values, *name = NULL, *base = NULL;
+    typecastObject *obj = NULL;
+
+    static char *kwlist[] = {"values", "name", "baseobj", NULL};
+
+    if (!PyArg_ParseTupleAndKeywords(args, keywds, "O!O!O!", kwlist,
+                                     &PyTuple_Type, &values,
+                                     &Text_Type, &name,
+                                     &typecastType, &base)) {
+        return NULL;
+    }
+
+    if ((obj = (typecastObject *)typecast_new(name, values, NULL, base))) {
+        obj->ccast = typecast_GENERIC_ARRAY_cast;
+        obj->pcast = NULL;
+    }
+
+    return (PyObject *)obj;
+}
+
 PyObject *
 typecast_from_c(typecastObject_initlist *type, PyObject *dict)
 {

psycopg/typecast.h

@@ -77,9 +77,11 @@ HIDDEN int typecast_add(PyObject *obj, PyObject *dict, int binary);
 /* the C callable typecastObject creator function */
 HIDDEN PyObject *typecast_from_c(typecastObject_initlist *type, PyObject *d);
 
-/* the python callable typecast creator function */
+/* the python callable typecast creator functions */
 HIDDEN PyObject *typecast_from_python(
     PyObject *self, PyObject *args, PyObject *keywds);
+HIDDEN PyObject *typecast_array_from_python(
+    PyObject *self, PyObject *args, PyObject *keywds);
 
 /* the function used to dispatch typecasting calls */
 HIDDEN PyObject *typecast_cast(

psycopg/typecast_array.c

@@ -133,7 +133,7 @@ typecast_array_tokenize(const char *str, Py_ssize_t strlength,
     }
 
     if (res == ASCAN_QUOTED) {
-        Py_ssize_t j;
+        const char *j, *jj;
         char *buffer = PyMem_Malloc(l+1);
         if (buffer == NULL) {
             PyErr_NoMemory();
@@ -142,10 +142,9 @@ typecast_array_tokenize(const char *str, Py_ssize_t strlength,
 
         *token = buffer;
 
-        for (j = *pos; j < *pos+l; j++) {
+        for (j = str + *pos, jj = j + l; j < jj; ++j) {
-            if (str[j] != '\\'
+            if (*j == '\\') { ++j; }
-                || (j > *pos && str[j-1] == '\\'))
+            *(buffer++) = *j;
-                *(buffer++) = str[j];
         }
 
         *buffer = '\0';

tests/types_basic.py

@@ -189,6 +189,17 @@ class TypesBasicTests(unittest.TestCase):
         s = self.execute("SELECT '{}'::text AS foo")
         self.failUnlessEqual(s, "{}")
 
+    def testArrayEscape(self):
+        ss = ['', '\\', '"', '\\\\', '\\"']
+        for s in ss:
+            r = self.execute("SELECT %s AS foo", (s,))
+            self.failUnlessEqual(s, r)
+            r = self.execute("SELECT %s AS foo", ([s],))
+            self.failUnlessEqual([s], r)
+
+        r = self.execute("SELECT %s AS foo", (ss,))
+        self.failUnlessEqual(ss, r)
+
     @testutils.skip_from_python(3)
     def testTypeRoundtripBuffer(self):
         o1 = buffer("".join(map(chr, range(256))))
@@ -285,6 +296,19 @@ class TypesBasicTests(unittest.TestCase):
         l1 = self.execute("select -%s;", (-1L,))
         self.assertEqual(1, l1)
 
+    def testGenericArray(self):
+        def caster(s, cur):
+            if s is None: return "nada"
+            return int(s) * 2
+        base = psycopg2.extensions.new_type((23,), "INT4", caster)
+        array = psycopg2.extensions.new_array_type((1007,), "INT4ARRAY", base)
+
+        psycopg2.extensions.register_type(array, self.conn)
+        a = self.execute("select '{1,2,3}'::int4[]")
+        self.assertEqual(a, [2,4,6])
+        a = self.execute("select '{1,2,NULL}'::int4[]")
+        self.assertEqual(a, [2,4,'nada'])
+
 
 class AdaptSubclassTest(unittest.TestCase):
     def test_adapt_subtype(self):

tests/types_extras.py

@@ -22,7 +22,7 @@ import re
 import sys
 from datetime import date
 
-from testutils import unittest, skip_if_no_uuid
+from testutils import unittest, skip_if_no_uuid, skip_before_postgres
 
 import psycopg2
 import psycopg2.extras
@@ -357,6 +357,63 @@ class HstoreTestCase(unittest.TestCase):
         finally:
             psycopg2.extensions.string_types.pop(oid)
 
+    @skip_if_no_hstore
+    @skip_before_postgres(8, 3)
+    def test_roundtrip_array(self):
+        from psycopg2.extras import register_hstore
+        register_hstore(self.conn)
+
+        ds = []
+        ds.append({})
+        ds.append({'a': 'b', 'c': None})
+
+        ab = map(chr, range(32, 128))
+        ds.append(dict(zip(ab, ab)))
+        ds.append({''.join(ab): ''.join(ab)})
+
+        self.conn.set_client_encoding('latin1')
+        if sys.version_info[0] < 3:
+            ab = map(chr, range(32, 127) + range(160, 255))
+        else:
+            ab = bytes(range(32, 127) + range(160, 255)).decode('latin1')
+
+        ds.append({''.join(ab): ''.join(ab)})
+        ds.append(dict(zip(ab, ab)))
+
+        cur = self.conn.cursor()
+        cur.execute("select %s", (ds,))
+        ds1 = cur.fetchone()[0]
+        self.assertEqual(ds, ds1)
+
+    @skip_if_no_hstore
+    @skip_before_postgres(8, 3)
+    def test_array_cast(self):
+        from psycopg2.extras import register_hstore
+        register_hstore(self.conn)
+        cur = self.conn.cursor()
+        cur.execute("select array['a=>1'::hstore, 'b=>2'::hstore];")
+        a = cur.fetchone()[0]
+        self.assertEqual(a, [{'a': '1'}, {'b': '2'}])
+
+    @skip_if_no_hstore
+    def test_array_cast_oid(self):
+        cur = self.conn.cursor()
+        cur.execute("select 'hstore'::regtype::oid, 'hstore[]'::regtype::oid")
+        oid, aoid = cur.fetchone()
+
+        from psycopg2.extras import register_hstore
+        register_hstore(None, globally=True, oid=oid, array_oid=aoid)
+        try:
+            cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore, '{a=>b}'::hstore[]")
+            t = cur.fetchone()
+            self.assert_(t[0] is None)
+            self.assertEqual(t[1], {})
+            self.assertEqual(t[2], {'a': 'b'})
+            self.assertEqual(t[3], [{'a': 'b'}])
+
+        finally:
+            psycopg2.extensions.string_types.pop(oid)
+            psycopg2.extensions.string_types.pop(aoid)
 
 def skip_if_no_composite(f):
     def skip_if_no_composite_(self):
@@ -563,6 +620,29 @@ class AdaptTypeTestCase(unittest.TestCase):
         curs.execute("select (4,8)::typens.typens_ii")
         self.assertEqual(curs.fetchone()[0], (4,8))
 
+    @skip_if_no_composite
+    @skip_before_postgres(8, 3)
+    def test_composite_array(self):
+        oid = self._create_type("type_isd",
+            [('anint', 'integer'), ('astring', 'text'), ('adate', 'date')])
+
+        t = psycopg2.extras.register_composite("type_isd", self.conn)
+
+        curs = self.conn.cursor()
+        r1 = (10, 'hello', date(2011,1,2))
+        r2 = (20, 'world', date(2011,1,3))
+        curs.execute("select %s::type_isd[];", ([r1, r2],))
+        v = curs.fetchone()[0]
+        self.assertEqual(len(v), 2)
+        self.assert_(isinstance(v[0], t.type))
+        self.assertEqual(v[0][0], 10)
+        self.assertEqual(v[0][1], "hello")
+        self.assertEqual(v[0][2], date(2011,1,2))
+        self.assert_(isinstance(v[1], t.type))
+        self.assertEqual(v[1][0], 20)
+        self.assertEqual(v[1][1], "world")
+        self.assertEqual(v[1][2], date(2011,1,3))
+
     def _create_type(self, name, fields):
         curs = self.conn.cursor()
         try: