diff --git a/NEWS b/NEWS index d8be8cb5..bb03987a 100644 --- a/NEWS +++ b/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. diff --git a/doc/src/advanced.rst b/doc/src/advanced.rst index 3ae15f05..4ccc3f15 100644 --- a/doc/src/advanced.rst +++ b/doc/src/advanced.rst @@ -216,6 +216,9 @@ read: >>> print type(point), point.x, point.y 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:: diff --git a/doc/src/extensions.rst b/doc/src/extensions.rst index 2fb1ae5c..ad933ef9 100644 --- a/doc/src/extensions.rst +++ b/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()`. diff --git a/lib/extensions.py b/lib/extensions.py index 10da66a0..de022ddf 100644 --- a/lib/extensions.py +++ b/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 diff --git a/lib/extras.py b/lib/extras.py index c6a2504f..78f962eb 100644 --- a/lib/extras.py +++ b/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 `) you may specify it in the - *oid* parameter (it can be found using a query such as :sql:`SELECT - 'hstore'::regtype::oid;`). + *oid* parameter, which can be found using a query such as :sql:`SELECT + '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 `). .. 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 diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c index f37a98e7..4dbf98d7 100644 --- a/psycopg/psycopgmodule.c +++ b/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}, diff --git a/psycopg/typecast.c b/psycopg/typecast.c index 56a203dc..8ede351c 100644 --- a/psycopg/typecast.c +++ b/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) { diff --git a/psycopg/typecast.h b/psycopg/typecast.h index cbae10a7..20def306 100644 --- a/psycopg/typecast.h +++ b/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( diff --git a/psycopg/typecast_array.c b/psycopg/typecast_array.c index 6bb256f1..75b84b50 100644 --- a/psycopg/typecast_array.c +++ b/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++) { - if (str[j] != '\\' - || (j > *pos && str[j-1] == '\\')) - *(buffer++) = str[j]; + for (j = str + *pos, jj = j + l; j < jj; ++j) { + if (*j == '\\') { ++j; } + *(buffer++) = *j; } *buffer = '\0'; diff --git a/tests/types_basic.py b/tests/types_basic.py index 709907eb..88c3844f 100755 --- a/tests/types_basic.py +++ b/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): diff --git a/tests/types_extras.py b/tests/types_extras.py index d6b5726f..c95a18c6 100755 --- a/tests/types_extras.py +++ b/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: