From abad3127ca0d658b6a4af10eb6017e4616e6ce43 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Sat, 6 Nov 2010 01:39:43 +0000 Subject: [PATCH] Added NamedTupleCursor. --- ChangeLog | 4 ++ NEWS-2.3 | 1 + doc/src/extras.rst | 64 +++++++++++++++++++---- lib/extras.py | 55 ++++++++++++++++++++ tests/extras_dictcursor.py | 103 +++++++++++++++++++++++++++++++++++++ 5 files changed, 217 insertions(+), 10 deletions(-) diff --git a/ChangeLog b/ChangeLog index 37e0aa29..47adb462 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,3 +1,7 @@ +2010-11-06 Daniele Varrazzo + + * lib/extras.py: added NamedTupleCursor. + 2010-11-05 Daniele Varrazzo * setup.py: bumped to version 2.3.dev0 diff --git a/NEWS-2.3 b/NEWS-2.3 index d8390d26..1b72d413 100644 --- a/NEWS-2.3 +++ b/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: diff --git a/doc/src/extras.rst b/doc/src/extras.rst index 9a988cf9..7e566d16 100644 --- a/doc/src/extras.rst +++ b/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 ----------------- diff --git a/lib/extras.py b/lib/extras.py index 902f4609..47dcbb23 100644 --- a/lib/extras.py +++ b/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. diff --git a/tests/extras_dictcursor.py b/tests/extras_dictcursor.py index 6c552c32..92028e23 100644 --- a/tests/extras_dictcursor.py +++ b/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__)