Added NamedTupleCursor.

ChangeLog

@@ -1,3 +1,7 @@
+2010-11-06  Daniele Varrazzo  <daniele.varrazzo@gmail.com>
+
+	* lib/extras.py: added NamedTupleCursor.
+
 2010-11-05  Daniele Varrazzo  <daniele.varrazzo@gmail.com>
 
 	* setup.py: bumped to version 2.3.dev0

NEWS-2.3

@@ -10,6 +10,7 @@ Psycopg 2.3 aims to expose some of the new features introduced in PostgreSQL
     9.0 and pre-9.0 syntax.
   - Two-phase commit protocol support as per DBAPI specification.
   - Support for payload in notifications received from the backed.
+  - namedtuple returning cursor.
 
 * Other changes:
 

doc/src/extras.rst

@@ -21,15 +21,23 @@ classes until a better place in the distribution is found.
 
 .. _dict-cursor:
 
+
+Connection and cursor subclasses
+--------------------------------
+
+A few objects that change the way the results are returned by the cursor or
+modify the object behavior in some other way. Typically `!connection`
+subclasses are passed as *connection_factory* argument to
+`~psycopg2.connect()` so that the connection will generate the matching
+`!cursor` subclass. Alternatively a `!cursor` subclass can be used one-off by
+passing it as the *cursor_factory* argument to the `~connection.cursor()`
+method of a regular `!connection`.
+
 Dictionary-like cursor
-----------------------
+^^^^^^^^^^^^^^^^^^^^^^
 
 The dict cursors allow to access to the retrieved records using an iterface
-similar to the Python dictionaries instead of the tuples. You can use it
-either passing `DictConnection` as `connection_factory` argument
-to the `~psycopg2.connect()` function or passing `DictCursor` as
-the `!cursor_factory` argument to the `~connection.cursor()` method
-of a regular `connection`.
+similar to the Python dictionaries instead of the tuples.
 
     >>> dict_cur = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
     >>> dict_cur.execute("INSERT INTO test (num, data) VALUES(%s, %s)",
@@ -67,11 +75,38 @@ Real dictionary cursor
 
 
 
+.. index::
+    pair: Cursor; namedtuple
+
+`namedtuple` cursor
+^^^^^^^^^^^^^^^^^^^^
+
+.. versionadded:: 2.3
+
+These objects require `!collection.namedtuple()` to be found, so it is
+available out-of-the-box only from Python 2.6. Anyway, the namedtuple
+implementation is compatible with previous Python versions, so all you
+have to do is to `download it`__ and add make it available where we
+expect it to be... ::
+
+    from somewhere import namedtuple
+    import collections
+    collections.namedtuple = namedtuple
+    from psycopg.extras import NamedTupleConnection
+    # ...
+
+.. __: http://code.activestate.com/recipes/500261-named-tuples/
+
+.. autoclass:: NamedTupleCursor
+
+.. autoclass:: NamedTupleConnection
+
+
 .. index::
     pair: Cursor; Logging
 
 Logging cursor
---------------
+^^^^^^^^^^^^^^
 
 .. autoclass:: LoggingConnection
     :members: initialize,filter
@@ -86,12 +121,19 @@ Logging cursor
 
 
 
+.. index::
+    single: Data types; Additional
+
+Additional data types
+---------------------
+
+
 .. index::
     pair: hstore; Data types
     pair: dict; Adaptation
 
 Hstore data type
-----------------
+^^^^^^^^^^^^^^^^
 
 .. versionadded:: 2.3
 
@@ -119,7 +161,7 @@ can be enabled using the `register_hstore()` function.
     pair: UUID; Data types
 
 UUID data type
---------------
+^^^^^^^^^^^^^^
 
 .. versionadded:: 2.0.9
 .. versionchanged:: 2.0.13 added UUID array support.
@@ -151,7 +193,7 @@ UUID data type
     pair: INET; Data types
 
 :sql:`inet` data type
-----------------------
+^^^^^^^^^^^^^^^^^^^^^^
 
 .. versionadded:: 2.0.9
 
@@ -190,6 +232,8 @@ Fractional time zones
 .. index::
    pair: Example; Coroutine;
 
+
+
 Coroutine support
 -----------------
 

lib/extras.py

@@ -234,6 +234,61 @@ class RealDictRow(dict):
         return dict.__setitem__(self, name, value)
 
 
+class NamedTupleConnection(_connection):
+    """A connection that uses `NamedTupleCursor` automatically."""
+    def cursor(self, *args, **kwargs):
+        kwargs['cursor_factory'] = NamedTupleCursor
+        return _connection.cursor(self, *args, **kwargs)
+
+class NamedTupleCursor(_cursor):
+    """A cursor that generates results as |namedtuple|__.
+
+    `!fetch*()` methods will return named tuples instead of regular tuples, so
+    their elements can be accessed both as regular numeric items as well as
+    attributes.
+
+        >>> nt_cur = conn.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor)
+        >>> rec = nt_cur.fetchone()
+        >>> rec
+        Record(id=1, num=100, data="abc'def")
+        >>> rec[1]
+        100
+        >>> rec.data
+        "abc'def"
+
+    .. |namedtuple| replace:: `!namedtuple`
+    .. __: http://docs.python.org/release/2.6/library/collections.html#collections.namedtuple
+    """
+    def fetchone(self):
+        t = _cursor.fetchone(self)
+        if t is not None:
+            nt = self._make_nt()
+            return nt(*t)
+
+    def fetchmany(self, size=None):
+        nt = self._make_nt()
+        ts = _cursor.fetchmany(self, size)
+        return [nt(*t) for t in ts]
+
+    def fetchall(self):
+        nt = self._make_nt()
+        ts = _cursor.fetchall(self)
+        return [nt(*t) for t in ts]
+
+    def __iter__(self):
+        return iter(self.fetchall())
+
+    try:
+        from collections import namedtuple
+    except ImportError, _exc:
+        def _make_nt(self):
+            raise self._exc
+    else:
+        def _make_nt(self, namedtuple=namedtuple):
+            return namedtuple("Record",
+                " ".join([d[0] for d in self.description]))
+
+
 class LoggingConnection(_connection):
     """A connection that logs all queries to a file or logger__ object.
 

tests/extras_dictcursor.py

@@ -105,6 +105,109 @@ class ExtrasDictCursorTests(unittest.TestCase):
         row = getter(curs)
         self.failUnless(row['foo'] == 'bar')
 
+
+def if_has_namedtuple(f):
+    def if_has_namedtuple_(self):
+        try:
+            from collections import namedtuple
+        except ImportError:
+            import warnings
+            warnings.warn("collections.namedtuple not available")
+        else:
+            return f(self)
+
+    if_has_namedtuple_.__name__ = f.__name__
+    return if_has_namedtuple_
+
+class NamedTupleCursorTest(unittest.TestCase):
+    def setUp(self):
+        from psycopg2.extras import NamedTupleConnection
+
+        try:
+            from collections import namedtuple
+        except ImportError:
+            self.conn = None
+            return
+
+        self.conn = psycopg2.connect(tests.dsn,
+            connection_factory=NamedTupleConnection)
+        curs = self.conn.cursor()
+        curs.execute("CREATE TEMPORARY TABLE nttest (i int, s text)")
+        curs.execute(
+            "INSERT INTO nttest VALUES (1, 'foo'), (2, 'bar'), (3, 'baz')")
+        self.conn.commit()
+
+    @if_has_namedtuple
+    def test_fetchone(self):
+        curs = self.conn.cursor()
+        curs.execute("select * from nttest where i = 1")
+        t = curs.fetchone()
+        self.assertEqual(t[0], 1)
+        self.assertEqual(t.i, 1)
+        self.assertEqual(t[1], 'foo')
+        self.assertEqual(t.s, 'foo')
+
+    @if_has_namedtuple
+    def test_fetchmany(self):
+        curs = self.conn.cursor()
+        curs.execute("select * from nttest order by 1")
+        res = curs.fetchmany(2)
+        self.assertEqual(2, len(res))
+        self.assertEqual(res[0].i, 1)
+        self.assertEqual(res[0].s, 'foo')
+        self.assertEqual(res[1].i, 2)
+        self.assertEqual(res[1].s, 'bar')
+
+    @if_has_namedtuple
+    def test_fetchall(self):
+        curs = self.conn.cursor()
+        curs.execute("select * from nttest order by 1")
+        res = curs.fetchall()
+        self.assertEqual(3, len(res))
+        self.assertEqual(res[0].i, 1)
+        self.assertEqual(res[0].s, 'foo')
+        self.assertEqual(res[1].i, 2)
+        self.assertEqual(res[1].s, 'bar')
+        self.assertEqual(res[2].i, 3)
+        self.assertEqual(res[2].s, 'baz')
+
+    @if_has_namedtuple
+    def test_iter(self):
+        curs = self.conn.cursor()
+        curs.execute("select * from nttest order by 1")
+        i = iter(curs)
+        t = i.next()
+        self.assertEqual(t.i, 1)
+        self.assertEqual(t.s, 'foo')
+        t = i.next()
+        self.assertEqual(t.i, 2)
+        self.assertEqual(t.s, 'bar')
+        t = i.next()
+        self.assertEqual(t.i, 3)
+        self.assertEqual(t.s, 'baz')
+        self.assertRaises(StopIteration, i.next)
+
+    def test_error_message(self):
+        try:
+            from collections import namedtuple
+        except ImportError:
+            # an import error somewhere
+            from psycopg2.extras import NamedTupleConnection
+            try:
+                self.conn = psycopg2.connect(tests.dsn,
+                    connection_factory=NamedTupleConnection)
+                curs = self.conn.cursor()
+                curs.execute("select 1")
+                curs.fetchone()
+            except ImportError:
+                pass
+            else:
+                self.fail("expecting ImportError")
+        else:
+            # skip the test
+            pass
+
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)