Merge branch 'contextmanager' into devel

NEWS

@@ -3,6 +3,8 @@ What's new in psycopg 2.5
 
   - Added JSON adaptation.
   - Added support for PostgreSQL 9.2 range types.
+  - 'connection' and 'cursor' objects can be used in 'with' statements
+    as context managers as specified by recent DBAPI extension.
   - Added support for backward scrollable cursors. Thanks to Jon Nelson
     for the initial patch (ticket #108).
   - Added a simple way to customize casting of composite types into Python

doc/src/connection.rst

@@ -74,6 +74,10 @@ The ``connection`` class
         automatically open, commands have immediate effect. See
         :ref:`transactions-control` for details.
 
+        .. versionchanged:: 2.5 if the connection is used in a ``with``
+            statement, the method is automatically called if no exception is
+            raised in the ``with`` block.
+
 
     .. index::
         pair: Transaction; Rollback
@@ -84,6 +88,10 @@ The ``connection`` class
         connection without committing the changes first will cause an implicit
         rollback to be performed.
 
+        .. versionchanged:: 2.5 if the connection is used in a ``with``
+            statement, the method is automatically called if an exception is
+            raised in the ``with`` block.
+
 
     .. method:: close()
 

doc/src/cursor.rst

@@ -83,6 +83,11 @@ The ``cursor`` class
         The cursor will be unusable from this point forward; an
         `~psycopg2.InterfaceError` will be raised if any operation is
         attempted with the cursor.
+
+        .. versionchanged:: 2.5 if the cursor is used in a ``with`` statement,
+            the method is automatically called at the end of the ``with``
+            block.
+
             
     .. attribute:: closed
 

doc/src/usage.rst

@@ -548,6 +548,30 @@ change the isolation level. See the `~connection.set_session()` method for all
 the details.
 
 
+.. index::
+    single: with statement
+
+``with`` statement
+^^^^^^^^^^^^^^^^^^
+
+Starting from version 2.5, psycopg2's connections and cursors are *context
+managers* and can be used with the ``with`` statement::
+
+    with psycopg2.connect(DSN) as conn:
+        with conn.cursor() as curs:
+            curs.execute(SQL)
+
+When a connection exits the ``with`` block, if no exception has been raised by
+the block, the transaction is committed. In case of exception the transaction
+is rolled back. In no case the connection is closed: a connection can be used
+in more than a ``with`` statement and each ``with`` block is effectively
+wrapped in a transaction.
+
+When a cursor exits the ``with`` block it is closed, releasing any resource
+eventually associated with it. The state of the transaction is not affected.
+
+
+
 .. index::
     pair: Server side; Cursor
     pair: Named; Cursor

psycopg/connection_type.c

@@ -124,7 +124,7 @@ exit:
 #define psyco_conn_close_doc "close() -- Close the connection."
 
 static PyObject *
-psyco_conn_close(connectionObject *self, PyObject *args)
+psyco_conn_close(connectionObject *self)
 {
     Dprintf("psyco_conn_close: closing connection at %p", self);
     conn_close(self);
@@ -140,7 +140,7 @@ psyco_conn_close(connectionObject *self, PyObject *args)
 #define psyco_conn_commit_doc "commit() -- Commit all changes to database."
 
 static PyObject *
-psyco_conn_commit(connectionObject *self, PyObject *args)
+psyco_conn_commit(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
     EXC_IF_CONN_ASYNC(self, commit);
@@ -160,7 +160,7 @@ psyco_conn_commit(connectionObject *self, PyObject *args)
 "rollback() -- Roll back all changes done to database."
 
 static PyObject *
-psyco_conn_rollback(connectionObject *self, PyObject *args)
+psyco_conn_rollback(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
     EXC_IF_CONN_ASYNC(self, rollback);
@@ -234,7 +234,7 @@ exit:
 "tpc_prepare() -- perform the first phase of a two-phase transaction."
 
 static PyObject *
-psyco_conn_tpc_prepare(connectionObject *self, PyObject *args)
+psyco_conn_tpc_prepare(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
     EXC_IF_CONN_ASYNC(self, tpc_prepare);
@@ -378,7 +378,7 @@ psyco_conn_tpc_rollback(connectionObject *self, PyObject *args)
 "tpc_recover() -- returns a list of pending transaction IDs."
 
 static PyObject *
-psyco_conn_tpc_recover(connectionObject *self, PyObject *args)
+psyco_conn_tpc_recover(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
     EXC_IF_CONN_ASYNC(self, tpc_recover);
@@ -389,6 +389,54 @@ psyco_conn_tpc_recover(connectionObject *self, PyObject *args)
 }
 
 
+#define psyco_conn_enter_doc \
+"__enter__ -> self"
+
+static PyObject *
+psyco_conn_enter(connectionObject *self)
+{
+    EXC_IF_CONN_CLOSED(self);
+
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+
+#define psyco_conn_exit_doc \
+"__exit__ -- commit if no exception, else roll back"
+
+static PyObject *
+psyco_conn_exit(connectionObject *self, PyObject *args)
+{
+    PyObject *type, *name, *tb;
+    PyObject *tmp = NULL;
+    PyObject *rv = NULL;
+
+    if (!PyArg_ParseTuple(args, "OOO", &type, &name, &tb)) {
+        goto exit;
+    }
+
+    if (type == Py_None) {
+        if (!(tmp = PyObject_CallMethod((PyObject *)self, "commit", ""))) {
+            goto exit;
+        }
+    } else {
+        if (!(tmp = PyObject_CallMethod((PyObject *)self, "rollback", ""))) {
+            goto exit;
+        }
+    }
+
+    /* success (of the commit or rollback, there may have been an exception in
+     * the block). Return None to avoid swallowing the exception */
+    rv = Py_None;
+    Py_INCREF(rv);
+
+exit:
+    Py_XDECREF(tmp);
+    return rv;
+}
+
+
 #ifdef PSYCOPG_EXTENSIONS
 
 
@@ -652,7 +700,7 @@ psyco_conn_set_client_encoding(connectionObject *self, PyObject *args)
 "get_transaction_status() -- Get backend transaction status."
 
 static PyObject *
-psyco_conn_get_transaction_status(connectionObject *self, PyObject *args)
+psyco_conn_get_transaction_status(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
 
@@ -756,7 +804,7 @@ psyco_conn_lobject(connectionObject *self, PyObject *args, PyObject *keywds)
 "get_backend_pid() -- Get backend process id."
 
 static PyObject *
-psyco_conn_get_backend_pid(connectionObject *self, PyObject *args)
+psyco_conn_get_backend_pid(connectionObject *self)
 {
     EXC_IF_CONN_CLOSED(self);
 
@@ -769,7 +817,7 @@ psyco_conn_get_backend_pid(connectionObject *self, PyObject *args)
 "reset() -- Reset current connection to defaults."
 
 static PyObject *
-psyco_conn_reset(connectionObject *self, PyObject *args)
+psyco_conn_reset(connectionObject *self)
 {
     int res;
 
@@ -797,7 +845,7 @@ psyco_conn_get_exception(PyObject *self, void *closure)
 }
 
 static PyObject *
-psyco_conn_poll(connectionObject *self, PyObject *args)
+psyco_conn_poll(connectionObject *self)
 {
     int res;
 
@@ -819,7 +867,7 @@ psyco_conn_poll(connectionObject *self, PyObject *args)
 "fileno() -> int -- Return file descriptor associated to database connection."
 
 static PyObject *
-psyco_conn_fileno(connectionObject *self, PyObject *args)
+psyco_conn_fileno(connectionObject *self)
 {
     long int socket;
 
@@ -838,7 +886,7 @@ psyco_conn_fileno(connectionObject *self, PyObject *args)
  "executing an asynchronous operation."
 
 static PyObject *
-psyco_conn_isexecuting(connectionObject *self, PyObject *args)
+psyco_conn_isexecuting(connectionObject *self)
 {
     /* synchronous connections will always return False */
     if (self->async == 0) {
@@ -870,7 +918,7 @@ psyco_conn_isexecuting(connectionObject *self, PyObject *args)
 "cancel() -- cancel the current operation"
 
 static PyObject *
-psyco_conn_cancel(connectionObject *self, PyObject *args)
+psyco_conn_cancel(connectionObject *self)
 {
     char errbuf[256];
 
@@ -924,6 +972,10 @@ static struct PyMethodDef connectionObject_methods[] = {
      METH_VARARGS, psyco_conn_tpc_rollback_doc},
     {"tpc_recover", (PyCFunction)psyco_conn_tpc_recover,
      METH_NOARGS, psyco_conn_tpc_recover_doc},
+    {"__enter__", (PyCFunction)psyco_conn_enter,
+     METH_NOARGS, psyco_conn_enter_doc},
+    {"__exit__", (PyCFunction)psyco_conn_exit,
+     METH_VARARGS, psyco_conn_exit_doc},
 #ifdef PSYCOPG_EXTENSIONS
     {"set_session", (PyCFunction)psyco_conn_set_session,
      METH_VARARGS|METH_KEYWORDS, psyco_conn_set_session_doc},

psycopg/cursor_type.c

@@ -50,7 +50,7 @@ extern PyObject *pyPsycopgTzFixedOffsetTimezone;
 "close() -- Close the cursor."
 
 static PyObject *
-psyco_curs_close(cursorObject *self, PyObject *args)
+psyco_curs_close(cursorObject *self)
 {
     EXC_IF_ASYNC_IN_PROGRESS(self, close);
 
@@ -761,7 +761,7 @@ exit:
 }
 
 static PyObject *
-psyco_curs_fetchone(cursorObject *self, PyObject *args)
+psyco_curs_fetchone(cursorObject *self)
 {
     PyObject *res;
 
@@ -953,7 +953,7 @@ exit:
 "Return `!None` when no more data is available.\n"
 
 static PyObject *
-psyco_curs_fetchall(cursorObject *self, PyObject *args)
+psyco_curs_fetchall(cursorObject *self)
 {
     int i, size;
     PyObject *list = NULL;
@@ -1085,7 +1085,7 @@ exit:
 "sets) and will raise a NotSupportedError exception."
 
 static PyObject *
-psyco_curs_nextset(cursorObject *self, PyObject *args)
+psyco_curs_nextset(cursorObject *self)
 {
     EXC_IF_CURS_CLOSED(self);
 
@@ -1201,6 +1201,42 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs)
 }
 
 
+#define psyco_curs_enter_doc \
+"__enter__ -> self"
+
+static PyObject *
+psyco_curs_enter(cursorObject *self)
+{
+    Py_INCREF(self);
+    return (PyObject *)self;
+}
+
+#define psyco_curs_exit_doc \
+"__exit__ -- close the cursor"
+
+static PyObject *
+psyco_curs_exit(cursorObject *self, PyObject *args)
+{
+    PyObject *tmp = NULL;
+    PyObject *rv = NULL;
+
+    /* don't care about the arguments here: don't need to parse them */
+
+    if (!(tmp = PyObject_CallMethod((PyObject *)self, "close", ""))) {
+        goto exit;
+    }
+
+    /* success (of curs.close()).
+     * Return None to avoid swallowing the exception */
+    rv = Py_None;
+    Py_INCREF(rv);
+
+exit:
+    Py_XDECREF(tmp);
+    return rv;
+}
+
+
 #ifdef PSYCOPG_EXTENSIONS
 
 /* Return a newly allocated buffer containing the list of columns to be
@@ -1674,7 +1710,7 @@ cursor_next(PyObject *self)
 
     if (NULL == ((cursorObject*)self)->name) {
         /* we don't parse arguments: psyco_curs_fetchone will do that for us */
-        res = psyco_curs_fetchone((cursorObject*)self, NULL);
+        res = psyco_curs_fetchone((cursorObject*)self);
 
         /* convert a None to NULL to signal the end of iteration */
         if (res && res == Py_None) {
@@ -1716,6 +1752,10 @@ static struct PyMethodDef cursorObject_methods[] = {
     /* DBAPI-2.0 extensions */
     {"scroll", (PyCFunction)psyco_curs_scroll,
      METH_VARARGS|METH_KEYWORDS, psyco_curs_scroll_doc},
+    {"__enter__", (PyCFunction)psyco_curs_enter,
+     METH_NOARGS, psyco_curs_enter_doc},
+    {"__exit__", (PyCFunction)psyco_curs_exit,
+     METH_VARARGS, psyco_curs_exit_doc},
     /* psycopg extensions */
 #ifdef PSYCOPG_EXTENSIONS
     {"cast", (PyCFunction)psyco_curs_cast,

tests/__init__.py

@@ -45,6 +45,11 @@ import test_transaction
 import test_types_basic
 import test_types_extras
 
+if sys.version_info[:2] >= (2, 5):
+    import test_with
+else:
+    test_with = None
+
 def test_suite():
     # If connection to test db fails, bail out early.
     import psycopg2
@@ -76,6 +81,8 @@ def test_suite():
     suite.addTest(test_transaction.test_suite())
     suite.addTest(test_types_basic.test_suite())
     suite.addTest(test_types_extras.test_suite())
+    if test_with:
+        suite.addTest(test_with.test_suite())
     return suite
 
 if __name__ == '__main__':

tests/test_with.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python
+
+# test_ctxman.py - unit test for connection and cursor used as context manager
+#
+# Copyright (C) 2012 Daniele Varrazzo  <daniele.varrazzo@gmail.com>
+#
+# psycopg2 is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# In addition, as a special exception, the copyright holders give
+# permission to link this program with the OpenSSL library (or with
+# modified versions of OpenSSL that use the same license as OpenSSL),
+# and distribute linked combinations including the two.
+#
+# You must obey the GNU Lesser General Public License in all respects for
+# all of the code used other than OpenSSL.
+#
+# psycopg2 is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
+# License for more details.
+
+
+from __future__ import with_statement
+
+import psycopg2
+import psycopg2.extensions as ext
+
+from testconfig import dsn
+from testutils import unittest
+
+class TestMixin(object):
+    def setUp(self):
+        self.conn = conn = psycopg2.connect(dsn)
+        curs = conn.cursor()
+        try:
+            curs.execute("delete from test_with")
+            conn.commit()
+        except psycopg2.ProgrammingError:
+            # assume table doesn't exist
+            conn.rollback()
+            curs.execute("create table test_with (id integer primary key)")
+            conn.commit()
+
+    def tearDown(self):
+        self.conn.close()
+
+
+class WithConnectionTestCase(TestMixin, unittest.TestCase):
+    def test_with_ok(self):
+        with self.conn as conn:
+            self.assert_(self.conn is conn)
+            self.assertEqual(conn.status, ext.STATUS_READY)
+            curs = conn.cursor()
+            curs.execute("insert into test_with values (1)")
+            self.assertEqual(conn.status, ext.STATUS_BEGIN)
+
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [(1,)])
+
+    def test_with_connect_idiom(self):
+        with psycopg2.connect(dsn) as conn:
+            self.assertEqual(conn.status, ext.STATUS_READY)
+            curs = conn.cursor()
+            curs.execute("insert into test_with values (2)")
+            self.assertEqual(conn.status, ext.STATUS_BEGIN)
+
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [(2,)])
+
+    def test_with_error_db(self):
+        def f():
+            with self.conn as conn:
+                curs = conn.cursor()
+                curs.execute("insert into test_with values ('a')")
+
+        self.assertRaises(psycopg2.DataError, f)
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [])
+
+    def test_with_error_python(self):
+        def f():
+            with self.conn as conn:
+                curs = conn.cursor()
+                curs.execute("insert into test_with values (3)")
+                1/0
+
+        self.assertRaises(ZeroDivisionError, f)
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [])
+
+    def test_with_closed(self):
+        def f():
+            with self.conn:
+                pass
+
+        self.conn.close()
+        self.assertRaises(psycopg2.InterfaceError, f)
+
+    def test_subclass_commit(self):
+        commits = []
+        class MyConn(ext.connection):
+            def commit(self):
+                commits.append(None)
+                super(MyConn, self).commit()
+
+        with psycopg2.connect(dsn, connection_factory=MyConn) as conn:
+            curs = conn.cursor()
+            curs.execute("insert into test_with values (10)")
+
+        self.assertEqual(conn.status, ext.STATUS_READY)
+        self.assert_(commits)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [(10,)])
+
+    def test_subclass_rollback(self):
+        rollbacks = []
+        class MyConn(ext.connection):
+            def rollback(self):
+                rollbacks.append(None)
+                super(MyConn, self).rollback()
+
+        try:
+            with psycopg2.connect(dsn, connection_factory=MyConn) as conn:
+                curs = conn.cursor()
+                curs.execute("insert into test_with values (11)")
+                1/0
+        except ZeroDivisionError:
+            pass
+        else:
+            self.assert_("exception not raised")
+
+        self.assertEqual(conn.status, ext.STATUS_READY)
+        self.assert_(rollbacks)
+
+        curs = conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [])
+
+
+class WithCursorTestCase(TestMixin, unittest.TestCase):
+    def test_with_ok(self):
+        with self.conn as conn:
+            with conn.cursor() as curs:
+                curs.execute("insert into test_with values (4)")
+                self.assert_(not curs.closed)
+            self.assertEqual(self.conn.status, ext.STATUS_BEGIN)
+            self.assert_(curs.closed)
+
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [(4,)])
+
+    def test_with_error(self):
+        try:
+            with self.conn as conn:
+                with conn.cursor() as curs:
+                    curs.execute("insert into test_with values (5)")
+                    1/0
+        except ZeroDivisionError:
+            pass
+
+        self.assertEqual(self.conn.status, ext.STATUS_READY)
+        self.assert_(not self.conn.closed)
+        self.assert_(curs.closed)
+
+        curs = self.conn.cursor()
+        curs.execute("select * from test_with")
+        self.assertEqual(curs.fetchall(), [])
+
+    def test_subclass(self):
+        closes = []
+        class MyCurs(ext.cursor):
+            def close(self):
+                closes.append(None)
+                super(MyCurs, self).close()
+
+        with self.conn.cursor(cursor_factory=MyCurs) as curs:
+            self.assert_(isinstance(curs, MyCurs))
+
+        self.assert_(curs.closed)
+        self.assert_(closes)
+
+
+def test_suite():
+    return unittest.TestLoader().loadTestsFromName(__name__)
+
+if __name__ == "__main__":
+    unittest.main()