From 0339042815dad143d60fc236ae9905cf2a137a3c Mon Sep 17 00:00:00 2001
From: Jason Kirtland <jek@discorporate.us>
Date: Thu, 16 Aug 2007 19:58:24 +0000
Subject: [PATCH] Added session.prune(), releases unused objects in strong-ref
 identity maps.

---
 CHANGES                          |  2 ++
 lib/sqlalchemy/orm/session.py    | 10 +++++++
 lib/sqlalchemy/orm/unitofwork.py | 21 ++++++++++++--
 test/orm/session.py              | 49 ++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/CHANGES b/CHANGES
index 1dcf0c42ba..277d162b7c 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,6 +1,8 @@
 0.4.0beta4
 - fix to bind param processing such that "False" values (like blank strings)
   still get processed/encoded
+- added session.prune(), trims away instances cached in a session that are
+  no longer referenced elsewhere. (a utility for for strong-ref identity maps)
 
 0.4.0beta3
 - sql types optimization:
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index c1e0e5b651..4fa9a4067a 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -725,6 +725,16 @@ class Session(object):
         for c in [obj] + list(_object_mapper(obj).cascade_iterator('refresh-expire', obj)):
             self._expire_impl(c)
 
+    def prune(self):
+        """Removes unreferenced instances cached in the identity map.
+
+        Removes any object in this Session'sidentity map that is not
+        referenced in user code, modified, new or scheduled for deletion.
+        Returns the number of objects pruned.
+        """
+        
+        return self.uow.prune_identity_map()
+
     def _expire_impl(self, obj):
         self._validate_persistent(obj)
 
diff --git a/lib/sqlalchemy/orm/unitofwork.py b/lib/sqlalchemy/orm/unitofwork.py
index 9cf3da639c..95ed950e19 100644
--- a/lib/sqlalchemy/orm/unitofwork.py
+++ b/lib/sqlalchemy/orm/unitofwork.py
@@ -19,12 +19,11 @@ new, dirty, or deleted and provides the capability to flush all those
 changes at once.
 """
 
+import gc, StringIO, weakref
 from sqlalchemy import util, logging, topological, exceptions
 from sqlalchemy.orm import attributes, interfaces
 from sqlalchemy.orm import util as mapperutil
 from sqlalchemy.orm.mapper import object_mapper
-import StringIO
-import weakref
 
 # Load lazily
 object_session = None
@@ -220,6 +219,24 @@ class UnitOfWork(object):
         if session.extension is not None:
             session.extension.after_flush_postexec(session, flush_context)
 
+    def prune_identity_map(self):
+        """Removes unreferenced instances cached in the identity map.
+
+        Removes any object in the identity map that is not referenced
+        in user code or scheduled for a unit of work operation.  Returns
+        the number of objects pruned.
+        """
+
+        if isinstance(self.identity_map, weakref.WeakValueDictionary):
+            return 0
+        ref_count = len(self.identity_map)
+        dirty = self.locate_dirty()
+        keepers = weakref.WeakValueDictionary(self.identity_map)
+        self.identity_map.clear()
+        gc.collect()
+        self.identity_map.update(keepers)
+        return ref_count - len(self.identity_map)
+
 class UOWTransaction(object):
     """Handles the details of organizing and executing transaction
     tasks during a UnitOfWork object's flush() operation.  
diff --git a/test/orm/session.py b/test/orm/session.py
index 1ad47dad6c..7857bf0647 100644
--- a/test/orm/session.py
+++ b/test/orm/session.py
@@ -385,6 +385,55 @@ class SessionTest(AssertMixin):
         import gc
         gc.collect()
         assert len(s.identity_map) == 1
+
+    def test_prune(self):
+        tables.delete()
+        s = create_session()
+        class User(object):pass
+        mapper(User, users)
+
+        for o in [User() for x in xrange(10)]:
+            s.save(o)
+        # o is still live after this loop...
+
+        self.assert_(len(s.identity_map) == 0)
+        self.assert_(s.prune() == 0)
+        s.flush()
+        self.assert_(s.prune() == 9)
+        self.assert_(len(s.identity_map) == 1)
+
+        user_id = o.user_id
+        del o
+        self.assert_(s.prune() == 1)
+        self.assert_(len(s.identity_map) == 0)
+
+        u = s.query(User).get(user_id)
+        self.assert_(s.prune() == 0)
+        self.assert_(len(s.identity_map) == 1)
+        u.user_name = 'squiznart'
+        del u
+        self.assert_(s.prune() == 0)
+        self.assert_(len(s.identity_map) == 1)
+        s.flush()
+        self.assert_(s.prune() == 1)
+        self.assert_(len(s.identity_map) == 0)
+
+        s.save(User())
+        self.assert_(s.prune() == 0)
+        self.assert_(len(s.identity_map) == 0)
+        s.flush()
+        self.assert_(len(s.identity_map) == 1)
+        self.assert_(s.prune() == 1)
+        self.assert_(len(s.identity_map) == 0)
+
+        u = s.query(User).get(user_id)
+        s.delete(u)
+        del u
+        self.assert_(s.prune() == 0)
+        self.assert_(len(s.identity_map) == 1)
+        s.flush()
+        self.assert_(s.prune() == 0)
+        self.assert_(len(s.identity_map) == 0)
         
     def test_no_save_cascade(self):
         mapper(Address, addresses)
-- 
2.47.3