Merge branch 'adapt-type' into python2

NEWS-2.3

@@ -1,6 +1,13 @@
 What's new in psycopg 2.3.3
 ---------------------------
 
+* New features:
+
+  - Added `register_composite()` function to cast PostgreSQL composite types
+    into Python tuples/namedtuples.
+
+* Bug fixes:
+
   - Fixed adaptation of None in composite types (ticket #26). Bug report by
     Karsten Hilbert.
 

doc/src/cursor.rst

@@ -122,20 +122,6 @@ The ``cursor`` class
         values can be retrieved using |fetch*|_ methods.
 
 
-    .. method:: mogrify(operation [, parameters])
-
-        Return a query string after arguments binding. The string returned is
-        exactly the one that would be sent to the database running the
-        `~cursor.execute()` method or similar.
-
-            >>> cur.mogrify("INSERT INTO test (num, data) VALUES (%s, %s)", (42, 'bar'))
-            "INSERT INTO test (num, data) VALUES (42, E'bar')"
-
-        .. extension::
-
-            The `mogrify()` method is a Psycopg extension to the |DBAPI|.
-
-        
     .. method:: executemany(operation, seq_of_parameters)
       
         Prepare a database operation (query or command) and then execute it
@@ -167,6 +153,34 @@ The ``cursor`` class
         does nothing but it is safe to call it.
 
 
+    .. method:: mogrify(operation [, parameters])
+
+        Return a query string after arguments binding. The string returned is
+        exactly the one that would be sent to the database running the
+        `~cursor.execute()` method or similar.
+
+            >>> cur.mogrify("INSERT INTO test (num, data) VALUES (%s, %s)", (42, 'bar'))
+            "INSERT INTO test (num, data) VALUES (42, E'bar')"
+
+        .. extension::
+
+            The `mogrify()` method is a Psycopg extension to the |DBAPI|.
+
+    .. method:: cast(oid, s)
+
+        Convert a value from the PostgreSQL string representation to a Python
+        object.
+
+        Use the most specific of the typecasters registered by
+        `~psycopg2.extensions.register_type()`.
+
+        .. versionadded:: 2.3.3
+
+        .. extension::
+
+            The `cast()` method is a Psycopg extension to the |DBAPI|.
+
+
 
     .. |fetch*| replace:: `!fetch*()`
 

doc/src/extras.rst

@@ -157,6 +157,55 @@ can be enabled using the `register_hstore()` function.
 
 
 
+.. index::
+    pair: Composite types; Data types
+    pair: tuple; Adaptation
+    pair: namedtuple; Adaptation
+
+Composite types casting
+^^^^^^^^^^^^^^^^^^^^^^^
+
+.. versionadded:: 2.3.3
+
+Using `register_composite()` it is possible to cast a PostgreSQL composite
+type (e.g. created with |CREATE TYPE|_ command) into a Python named tuple, or
+into a regular tuple if `!collections.namedtuple()` is not found.
+
+.. |CREATE TYPE| replace:: :sql:`CREATE TYPE`
+.. _CREATE TYPE: http://www.postgresql.org/docs/9.0/static/sql-createtype.html
+
+.. doctest::
+
+    >>> cur.execute("CREATE TYPE card AS (value int, suit text);")
+    >>> psycopg2.extras.register_composite('card', cur)
+    <psycopg2.extras.CompositeCaster object at 0x...>
+
+    >>> cur.execute("select (8, 'hearts')::card")
+    >>> cur.fetchone()[0]
+    card(value=8, suit='hearts')
+
+Nested composite types are handled as expected, but the type of the composite
+components must be registered as well.
+
+.. doctest::
+
+    >>> cur.execute("CREATE TYPE card_back AS (face card, back text);")
+    >>> psycopg2.extras.register_composite('card_back', cur)
+    <psycopg2.extras.CompositeCaster object at 0x...>
+
+    >>> cur.execute("select ((8, 'hearts'), 'blue')::card_back")
+    >>> cur.fetchone()[0]
+    card_back(face=card(value=8, suit='hearts'), back='blue')
+
+Adaptation from Python tuples to composite types is automatic instead and
+requires no adapter registration.
+
+.. autofunction:: register_composite
+
+.. autoclass:: CompositeCaster
+
+
+
 .. index::
     pair: UUID; Data types
 

doc/src/usage.rst

@@ -280,7 +280,7 @@ the SQL string that would be sent to the database.
     single: IN operator
 
 - Python tuples are converted in a syntax suitable for the SQL :sql:`IN`
-  operator::
+  operator and to represent a composite type::
 
     >>> cur.mogrify("SELECT %s IN %s;", (10, (10, 20, 30)))
     'SELECT 10 IN (10, 20, 30);'
@@ -290,6 +290,10 @@ the SQL string that would be sent to the database.
     SQL doesn't allow an empty list in the IN operator, so your code should
     guard against empty tuples.
 
+  If you want PostgreSQL composite types to be converted into a Python
+  tuple/namedtuple you can use the `~psycopg2.extras.register_composite()`
+  function.
+
   .. versionadded:: 2.0.6
      the tuple :sql:`IN` adaptation.
 
@@ -298,6 +302,10 @@ the SQL string that would be sent to the database.
      was necessary to import the `~psycopg2.extensions` module to have it
      registered.
 
+  .. versionchanged:: 2.3
+     named tuples are adapted like regular tuples and can thus be used to
+     represent composite types.
+
 - Python dictionaries are converted into the |hstore|_ data type. See
   `~psycopg2.extras.register_hstore()` for further details.
 
@@ -307,6 +315,7 @@ the SQL string that would be sent to the database.
   .. versionadded:: 2.3
      the :sql:`hstore` adaptation.
 
+
 .. index::
     single: Unicode
 

lib/extras.py

@@ -731,4 +731,146 @@ def register_hstore(conn_or_curs, globally=False, unicode=False):
     _ext.register_adapter(dict, HstoreAdapter)
 
 
+class CompositeCaster(object):
+    """Helps conversion of a PostgreSQL composite type into a Python object.
+
+    The class is usually created by the `register_composite()` function.
+
+    .. attribute:: name
+
+        The name of the PostgreSQL type.
+
+    .. attribute:: oid
+
+        The oid of the PostgreSQL type.
+
+    .. attribute:: type
+
+        The type of the Python objects returned. If `!collections.namedtuple()`
+        is available, it is a named tuple with attributes equal to the type
+        components. Otherwise it is just the `tuple` object.
+
+    .. attribute:: attnames
+
+        List of component names of the type to be casted.
+
+    .. attribute:: atttypes
+
+        List of component type oids of the type to be casted.
+
+    """
+    def __init__(self, name, oid, attrs):
+        self.name = name
+        self.oid = oid
+
+        self.attnames = [ a[0] for a in attrs ]
+        self.atttypes = [ a[1] for a in attrs ]
+        self.type = self._create_type(name, self.attnames)
+        self.typecaster = _ext.new_type((oid,), name, self.parse)
+
+    def parse(self, s, curs):
+        if s is None:
+            return None
+
+        tokens = self.tokenize(s)
+        if len(tokens) != len(self.atttypes):
+            raise psycopg2.DataError(
+                "expecting %d components for the type %s, %d found instead",
+                (len(self.atttypes), self.name, len(self.tokens)))
+
+        attrs = [ curs.cast(oid, token)
+            for oid, token in zip(self.atttypes, tokens) ]
+        return self.type(*attrs)
+
+    _re_tokenize = regex.compile(r"""
+  \(? ([,\)])                       # an empty token, representing NULL
+| \(? " ((?: [^"] | "")*) " [,)]    # or a quoted string
+| \(? ([^",\)]+) [,\)]              # or an unquoted string
+    """, regex.VERBOSE)
+
+    _re_undouble = regex.compile(r'(["\\])\1')
+
+    @classmethod
+    def tokenize(self, s):
+        rv = []
+        for m in self._re_tokenize.finditer(s):
+            if m is None:
+                raise psycopg2.InterfaceError("can't parse type: %r", s)
+            if m.group(1):
+                rv.append(None)
+            elif m.group(2):
+                rv.append(self._re_undouble.sub(r"\1", m.group(2)))
+            else:
+                rv.append(m.group(3))
+
+        return rv
+
+    def _create_type(self, name, attnames):
+        try:
+            from collections import namedtuple
+        except ImportError:
+            return tuple
+        else:
+            return namedtuple(name, attnames)
+
+    @classmethod
+    def _from_db(self, name, conn_or_curs):
+        """Return a `CompositeCaster` instance for the type *name*.
+
+        Raise `ProgrammingError` if the type is not found.
+        """
+        if hasattr(conn_or_curs, 'execute'):
+            conn = conn_or_curs.connection
+            curs = conn_or_curs
+        else:
+            conn = conn_or_curs
+            curs = conn_or_curs.cursor()
+
+        # Store the transaction status of the connection to revert it after use
+        conn_status = conn.status
+
+        # get the type oid and attributes
+        curs.execute("""\
+SELECT t.oid, 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 = 'public';
+""", (name, ))
+
+        recs = curs.fetchall()
+
+        # revert the status of the connection as before the command
+        if (conn_status != _ext.STATUS_IN_TRANSACTION
+        and conn.isolation_level != _ext.ISOLATION_LEVEL_AUTOCOMMIT):
+            conn.rollback()
+
+        if not recs:
+            raise psycopg2.ProgrammingError(
+                "PostgreSQL type '%s' not found" % name)
+
+        type_oid = recs[0][0]
+        type_attrs = [ (r[1], r[2]) for r in recs ]
+
+        return CompositeCaster(name, type_oid, type_attrs)
+
+def register_composite(name, conn_or_curs, globally=False):
+    """Register a typecaster to convert a composite type into a tuple.
+
+    :param name: the name of a PostgreSQL composite type, e.g. created using
+        the |CREATE TYPE|_ command
+    :param conn_or_curs: a connection or cursor used to find the type oid and
+        components; the typecaster is registered in a scope limited to this
+        object, unless *globally* is set to `True`
+    :param globally: if `False` (default) register the typecaster only on
+        *conn_or_curs*, otherwise register it globally
+    :return: the registered `CompositeCaster` instance responsible for the
+        conversion
+    """
+    caster = CompositeCaster._from_db(name, conn_or_curs)
+    _ext.register_type(caster.typecaster, not globally and conn_or_curs or None)
+
+    return caster
+
+
 __all__ = filter(lambda k: not k.startswith('_'), locals().keys())

psycopg/cursor.h

@@ -84,6 +84,7 @@ typedef struct {
 } cursorObject;
 
 /* C-callable functions in cursor_int.c and cursor_ext.c */
+HIDDEN PyObject *curs_get_cast(cursorObject *self, PyObject *oid);
 HIDDEN void curs_reset(cursorObject *self);
 
 /* exception-raising macros */

psycopg/cursor_int.c

@@ -32,6 +32,41 @@
 #include "psycopg/psycopg.h"
 #include "psycopg/cursor.h"
 #include "psycopg/pqpath.h"
+#include "psycopg/typecast.h"
+
+/* curs_get_cast - return the type caster for an oid.
+ *
+ * Return the most specific type caster, from cursor to connection to global.
+ * If no type caster is found, return the default one.
+ *
+ * Return a borrowed reference.
+ */
+
+PyObject *
+curs_get_cast(cursorObject *self, PyObject *oid)
+{
+    PyObject *cast;
+
+    /* cursor lookup */
+    if (self->string_types != NULL && self->string_types != Py_None) {
+        cast = PyDict_GetItem(self->string_types, oid);
+        Dprintf("curs_get_cast:        per-cursor dict: %p", cast);
+        if (cast) { return cast; }
+    }
+
+    /* connection lookup */
+    cast = PyDict_GetItem(self->conn->string_types, oid);
+    Dprintf("curs_get_cast:        per-connection dict: %p", cast);
+    if (cast) { return cast; }
+
+    /* global lookup */
+    cast = PyDict_GetItem(psyco_types, oid);
+    Dprintf("curs_get_cast:        global dict: %p", cast);
+    if (cast) { return cast; }
+
+    /* fallback */
+    return psyco_default_cast;
+}
 
 /* curs_reset - reset the cursor to a clean state */
 

psycopg/cursor_type.c

@@ -629,6 +629,29 @@ psyco_curs_mogrify(cursorObject *self, PyObject *args, PyObject *kwargs)
 #endif
 
 
+/* cast method - convert an oid/string into a Python object */
+#define psyco_curs_cast_doc \
+"cast(oid, s) -> value\n\n" \
+"Convert the string s to a Python object according to its oid.\n\n" \
+"Look for a typecaster first in the cursor, then in its connection," \
+"then in the global register. If no suitable typecaster is found," \
+"leave the value as a string."
+
+static PyObject *
+psyco_curs_cast(cursorObject *self, PyObject *args)
+{
+    PyObject *oid;
+    PyObject *s;
+    PyObject *cast;
+
+    if (!PyArg_ParseTuple(args, "OO", &oid, &s))
+        return NULL;
+
+    cast = curs_get_cast(self, oid);
+    return PyObject_CallFunctionObjArgs(cast, s, (PyObject *)self, NULL);
+}
+
+
 /* fetchone method - fetch one row of results */
 
 #define psyco_curs_fetchone_doc \
@@ -1524,6 +1547,8 @@ static struct PyMethodDef cursorObject_methods[] = {
      METH_VARARGS|METH_KEYWORDS, psyco_curs_scroll_doc},
     /* psycopg extensions */
 #ifdef PSYCOPG_EXTENSIONS
+    {"cast", (PyCFunction)psyco_curs_cast,
+     METH_VARARGS, psyco_curs_cast_doc},
     {"mogrify", (PyCFunction)psyco_curs_mogrify,
      METH_VARARGS|METH_KEYWORDS, psyco_curs_mogrify_doc},
     {"copy_from", (PyCFunction)psyco_curs_copy_from,

psycopg/pqpath.c

@@ -958,19 +958,7 @@ _pq_fetch_tuples(cursorObject *curs)
 
         type = PyInt_FromLong(ftype);
         Dprintf("_pq_fetch_tuples: looking for cast %d:", ftype);
-        if (curs->string_types != NULL && curs->string_types != Py_None) {
-            cast = PyDict_GetItem(curs->string_types, type);
-            Dprintf("_pq_fetch_tuples:     per-cursor dict: %p", cast);
-        }
-        if (cast == NULL) {
-            cast = PyDict_GetItem(curs->conn->string_types, type);
-            Dprintf("_pq_fetch_tuples:     per-connection dict: %p", cast);
-        }
-        if (cast == NULL) {
-            cast = PyDict_GetItem(psyco_types, type);
-            Dprintf("_pq_fetch_tuples:     global dict: %p", cast);
-        }
-        if (cast == NULL) cast = psyco_default_cast;
+        cast = curs_get_cast(curs, type);
 
         /* else if we got binary tuples and if we got a field that
            is binary use the default cast

tests/test_cursor.py

@@ -67,6 +67,38 @@ class CursorTests(unittest.TestCase):
         self.assertEqual('SELECT 10.3;',
             cur.mogrify("SELECT %s;", (Decimal("10.3"),)))
 
+    def test_cast(self):
+        curs = self.conn.cursor()
+
+        self.assertEqual(42, curs.cast(20, '42'))
+        self.assertAlmostEqual(3.14, curs.cast(700, '3.14'))
+
+        try:
+            from decimal import Decimal
+        except ImportError:
+            self.assertAlmostEqual(123.45, curs.cast(1700, '123.45'))
+        else:
+            self.assertEqual(Decimal('123.45'), curs.cast(1700, '123.45'))
+
+        from datetime import date
+        self.assertEqual(date(2011,1,2), curs.cast(1082, '2011-01-02'))
+        self.assertEqual("who am i?", curs.cast(705, 'who am i?'))  # unknown
+
+    def test_cast_specificity(self):
+        curs = self.conn.cursor()
+        self.assertEqual("foo", curs.cast(705, 'foo'))
+
+        D = psycopg2.extensions.new_type((705,), "DOUBLING", lambda v, c: v * 2)
+        psycopg2.extensions.register_type(D, self.conn)
+        self.assertEqual("foofoo", curs.cast(705, 'foo'))
+
+        T = psycopg2.extensions.new_type((705,), "TREBLING", lambda v, c: v * 3)
+        psycopg2.extensions.register_type(T, curs)
+        self.assertEqual("foofoofoo", curs.cast(705, 'foo'))
+
+        curs2 = self.conn.cursor()
+        self.assertEqual("foofoo", curs2.cast(705, 'foo'))
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)

tests/types_extras.py

@@ -20,6 +20,8 @@ except:
     pass
 import re
 import sys
+from datetime import date
+
 from testutils import unittest
 
 import psycopg2
@@ -385,6 +387,158 @@ class AdaptTypeTestCase(unittest.TestCase):
         finally:
             ext.register_adapter(type(None), orig_adapter)
 
+    def test_tokenization(self):
+        from psycopg2.extras import CompositeCaster
+        def ok(s, v):
+            self.assertEqual(CompositeCaster.tokenize(s), v)
+
+        ok("(,)", [None, None])
+        ok('(hello,,10.234,2010-11-11)', ['hello', None, '10.234', '2010-11-11'])
+        ok('(10,"""")', ['10', '"'])
+        ok('(10,",")', ['10', ','])
+        ok(r'(10,"\\")', ['10', '\\'])
+        ok(r'''(10,"\\',""")''', ['10', '''\\',"'''])
+        ok('(10,"(20,""(30,40)"")")', ['10', '(20,"(30,40)")'])
+        ok('(10,"(20,""(30,""""(40,50)"""")"")")', ['10', '(20,"(30,""(40,50)"")")'])
+        ok('(,"(,""(a\nb\tc)"")")', [None, '(,"(a\nb\tc)")'])
+        ok('(\x01,\x02,\x03,\x04,\x05,\x06,\x07,\x08,"\t","\n","\x0b",'
+           '"\x0c","\r",\x0e,\x0f,\x10,\x11,\x12,\x13,\x14,\x15,\x16,'
+           '\x17,\x18,\x19,\x1a,\x1b,\x1c,\x1d,\x1e,\x1f," ",!,"""",#,'
+           '$,%,&,\',"(",")",*,+,",",-,.,/,0,1,2,3,4,5,6,7,8,9,:,;,<,=,>,?,'
+           '@,A,B,C,D,E,F,G,H,I,J,K,L,M,N,O,P,Q,R,S,T,U,V,W,X,Y,Z,[,"\\\\",],'
+           '^,_,`,a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,{,|,},'
+           '~,\x7f)',
+           map(chr, range(1, 128)))
+        ok('(,"\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f'
+           '\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f !'
+           '""#$%&\'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\\\]'
+           '^_`abcdefghijklmnopqrstuvwxyz{|}~\x7f")',
+           [None, ''.join(map(chr, range(1, 128)))])
+
+    def test_cast_composite(self):
+        oid = self._create_type("type_isd",
+            [('anint', 'integer'), ('astring', 'text'), ('adate', 'date')])
+
+        t = psycopg2.extras.register_composite("type_isd", self.conn)
+        self.assertEqual(t.name, 'type_isd')
+        self.assertEqual(t.oid, oid)
+        self.assert_(issubclass(t.type, tuple))
+        self.assertEqual(t.attnames, ['anint', 'astring', 'adate'])
+        self.assertEqual(t.atttypes, [23,25,1082])
+
+        curs = self.conn.cursor()
+        r = (10, 'hello', date(2011,1,2))
+        curs.execute("select %s::type_isd;", (r,))
+        v = curs.fetchone()[0]
+        self.assert_(isinstance(v, t.type))
+        self.assertEqual(v[0], 10)
+        self.assertEqual(v[1], "hello")
+        self.assertEqual(v[2], date(2011,1,2))
+
+        try:
+            from collections import namedtuple
+        except ImportError:
+            pass
+        else:
+            self.assert_(t.type is not tuple)
+            self.assertEqual(v.anint, 10)
+            self.assertEqual(v.astring, "hello")
+            self.assertEqual(v.adate, date(2011,1,2))
+
+    def test_cast_nested(self):
+        self._create_type("type_is",
+            [("anint", "integer"), ("astring", "text")])
+        self._create_type("type_r_dt",
+            [("adate", "date"), ("apair", "type_is")])
+        self._create_type("type_r_ft",
+            [("afloat", "float8"), ("anotherpair", "type_r_dt")])
+
+        psycopg2.extras.register_composite("type_is", self.conn)
+        psycopg2.extras.register_composite("type_r_dt", self.conn)
+        psycopg2.extras.register_composite("type_r_ft", self.conn)
+
+        curs = self.conn.cursor()
+        r = (0.25, (date(2011,1,2), (42, "hello")))
+        curs.execute("select %s::type_r_ft;", (r,))
+        v = curs.fetchone()[0]
+
+        self.assertEqual(r, v)
+
+        try:
+            from collections import namedtuple
+        except ImportError:
+            pass
+        else:
+            self.assertEqual(v.anotherpair.apair.astring, "hello")
+
+    def test_register_on_cursor(self):
+        self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
+
+        curs1 = self.conn.cursor()
+        curs2 = self.conn.cursor()
+        psycopg2.extras.register_composite("type_ii", curs1)
+        curs1.execute("select (1,2)::type_ii")
+        self.assertEqual(curs1.fetchone()[0], (1,2))
+        curs2.execute("select (1,2)::type_ii")
+        self.assertEqual(curs2.fetchone()[0], "(1,2)")
+
+    def test_register_on_connection(self):
+        self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
+
+        conn1 = psycopg2.connect(self.conn.dsn)
+        conn2 = psycopg2.connect(self.conn.dsn)
+        try:
+            psycopg2.extras.register_composite("type_ii", conn1)
+            curs1 = conn1.cursor()
+            curs2 = conn2.cursor()
+            curs1.execute("select (1,2)::type_ii")
+            self.assertEqual(curs1.fetchone()[0], (1,2))
+            curs2.execute("select (1,2)::type_ii")
+            self.assertEqual(curs2.fetchone()[0], "(1,2)")
+        finally:
+            conn1.close()
+            conn2.close()
+
+    def test_register_globally(self):
+        self._create_type("type_ii", [("a", "integer"), ("b", "integer")])
+
+        conn1 = psycopg2.connect(self.conn.dsn)
+        conn2 = psycopg2.connect(self.conn.dsn)
+        try:
+            t = psycopg2.extras.register_composite("type_ii", conn1, globally=True)
+            try:
+                curs1 = conn1.cursor()
+                curs2 = conn2.cursor()
+                curs1.execute("select (1,2)::type_ii")
+                self.assertEqual(curs1.fetchone()[0], (1,2))
+                curs2.execute("select (1,2)::type_ii")
+                self.assertEqual(curs2.fetchone()[0], (1,2))
+            finally:
+                del psycopg2.extensions.string_types[t.oid]
+
+        finally:
+            conn1.close()
+            conn2.close()
+
+    def _create_type(self, name, fields):
+        curs = self.conn.cursor()
+        try:
+            curs.execute("drop type %s cascade;" % name)
+        except psycopg2.ProgrammingError:
+            self.conn.rollback()
+
+        curs.execute("create type %s as (%s);" % (name,
+            ", ".join(["%s %s" % p for p in fields])))
+        curs.execute("""\
+            SELECT t.oid
+            FROM pg_type t JOIN pg_namespace ns ON typnamespace = ns.oid
+            WHERE typname = %s and nspname = 'public';
+            """, (name,))
+        oid = curs.fetchone()[0]
+        self.conn.commit()
+        return oid
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)