Added 'get_wait_callback()' function.

4 files changed, 29 insertions, 2 deletions

diff --git a/lib/extensions.py b/lib/extensions.py
index f7431d1..7290d5b 100644
--- a/lib/extensions.py
+++ b/lib/extensions.py
@@ -61,7 +61,7 @@ from _psycopg import ISQLQuote
 from _psycopg import QueryCanceledError, TransactionRollbackError
 
 try:
-    from _psycopg import set_wait_callback
+    from _psycopg import set_wait_callback, get_wait_callback
 except ImportError:
     pass
 
diff --git a/psycopg/green.c b/psycopg/green.c
index a49f818..2a66d94 100644
--- a/psycopg/green.c
+++ b/psycopg/green.c
@@ -56,6 +56,25 @@ psyco_set_wait_callback(PyObject *self, PyObject *obj)
 }
 
 
+/* Return the currently registered wait callback function.
+ *
+ * The function is exported by the _psycopg module.
+ */
+PyObject *
+psyco_get_wait_callback(PyObject *self, PyObject *obj)
+{
+    PyObject *ret;
+
+    ret = wait_callback;
+    if (!ret) {
+        ret = Py_None;
+    }
+
+    Py_INCREF(ret);
+    return ret;
+}
+
+
 /* Return nonzero if a wait callback should be called. */
 int
 psyco_green()
diff --git a/psycopg/green.h b/psycopg/green.h
index 3c09b32..a21a9a1 100644
--- a/psycopg/green.h
+++ b/psycopg/green.h
@@ -47,8 +47,14 @@ extern "C" {
 "\n" \
 "See `~psycopg2.extras.wait_select()` for an example of a wait callback\n" \
 "implementation.\n"
-
 HIDDEN PyObject *psyco_set_wait_callback(PyObject *self, PyObject *obj);
+
+#define psyco_get_wait_callback_doc \
+"get_wait_callback() -- Return the currently registered wait callback.\n" \
+"\n" \
+"Return `None` if no callback is currently registered.\n"
+HIDDEN PyObject *psyco_get_wait_callback(PyObject *self, PyObject *obj);
+
 HIDDEN int psyco_green(void);
 HIDDEN PyObject *psyco_wait(PyObject *conn, PyObject *curs);
 HIDDEN PGresult *psyco_exec_green(connectionObject *conn, const char *command);
diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c
index 60c9c04..8b47944 100644
--- a/psycopg/psycopgmodule.c
+++ b/psycopg/psycopgmodule.c
@@ -697,6 +697,8 @@ static PyMethodDef psycopgMethods[] = {
 #ifdef PSYCOPG_EXTENSIONS
     {"set_wait_callback",  (PyCFunction)psyco_set_wait_callback,
      METH_O, psyco_set_wait_callback_doc},
+    {"get_wait_callback",  (PyCFunction)psyco_get_wait_callback,
+     METH_NOARGS, psyco_get_wait_callback_doc},
 #endif
 
     {NULL, NULL, 0, NULL}        /* Sentinel */