From 04cf90cc21db4a8213cfc07da3a2920c8d178faa Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 3 Jan 2011 11:52:38 +0100
Subject: [PATCH] The connection is weakly referenceable

---
 NEWS-2.3                  |  1 +
 psycopg/connection.h      |  1 +
 psycopg/connection_type.c | 15 +++++++++++----
 tests/test_connection.py  |  8 ++++++++
 4 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/NEWS-2.3 b/NEWS-2.3
index fde7bd65..82b87594 100644
--- a/NEWS-2.3
+++ b/NEWS-2.3
@@ -6,6 +6,7 @@ What's new in psycopg 2.3.3
   - Added `register_composite()` function to cast PostgreSQL composite types
     into Python tuples/namedtuples.
   - The build script refuses to guess values if pg_config is not found.
+  - Connections are weakly referenceable.
 
 * Bug fixes:
 
diff --git a/psycopg/connection.h b/psycopg/connection.h
index 76a6a093..72437492 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -119,6 +119,7 @@ typedef struct {
     PyObject *binary_types;   /* a set of typecasters for binary types */
 
     int equote;               /* use E''-style quotes for escaped strings */
+    PyObject *weakreflist;    /* list of weak references */
 
 } connectionObject;
 
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 3469dc2f..15d943a2 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -867,6 +867,7 @@ connection_setup(connectionObject *self, const char *dsn, long int async)
     self->binary_types = PyDict_New();
     self->notice_pending = NULL;
     self->encoding = NULL;
+    self->weakreflist = NULL;
 
     pthread_mutex_init(&(self->lock), NULL);
 
@@ -896,11 +897,15 @@ static void
 connection_dealloc(PyObject* obj)
 {
     connectionObject *self = (connectionObject *)obj;
-    
+
+    if (self->weakreflist) {
+        PyObject_ClearWeakRefs(obj);
+    }
+
     PyObject_GC_UnTrack(self);
 
     if (self->closed == 0) conn_close(self);
-    
+
     conn_notice_clean(self);
 
     if (self->dsn) free(self->dsn);
@@ -1002,14 +1007,16 @@ PyTypeObject connectionType = {
     0,          /*tp_setattro*/
     0,          /*tp_as_buffer*/
 
-    Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE|Py_TPFLAGS_HAVE_GC, /*tp_flags*/
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC |
+        Py_TPFLAGS_HAVE_WEAKREFS,
+                /*tp_flags*/
     connectionType_doc, /*tp_doc*/
 
     (traverseproc)connection_traverse, /*tp_traverse*/
     0,          /*tp_clear*/
 
     0,          /*tp_richcompare*/
-    0,          /*tp_weaklistoffset*/
+    offsetof(connectionObject, weakreflist), /* tp_weaklistoffset */
 
     0,          /*tp_iter*/
     0,          /*tp_iternext*/
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 8d87e730..396254f2 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -111,6 +111,14 @@ class ConnectionTests(unittest.TestCase):
         self.assert_(time.time() - t0 < 3,
             "something broken in concurrency")
 
+    def test_weakref(self):
+        from weakref import ref
+        conn = psycopg2.connect(self.conn.dsn)
+        w = ref(conn)
+        conn.close()
+        del conn
+        self.assert_(w() is None)
+
 
 class IsolationLevelsTestCase(unittest.TestCase):