Rev 66: Add a 'scanner.get_recursive_size' function. in http://bazaar.launchpad.net/~jameinel/meliae/trunk

[John Arbash Meinel]

At http://bazaar.launchpad.net/~jameinel/meliae/trunk

------------------------------------------------------------
revno: 66
revision-id: john at arbash-meinel.com-20090908174103-17nem9vf7bz8r2x7
parent: john at arbash-meinel.com-20090908171916-ekr6d1c8gb0e0eih
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: trunk
timestamp: Tue 2009-09-08 12:41:03 -0500
message:
  Add a 'scanner.get_recursive_size' function.
  This is meant to help in interactive analysis. So that we
  can try to quickly determine the size of a simple structure
  in memory.
-------------- next part --------------
=== modified file 'meliae/scanner.py'
--- a/meliae/scanner.py	2009-09-08 17:19:16 +0000
+++ b/meliae/scanner.py	2009-09-08 17:41:03 +0000
@@ -80,3 +80,35 @@
     for obj in all_objs:
         _scanner.dump_object_info(outf, obj, nodump=nodump,
                                   recurse_depth=recurse_depth)
+
+
+def get_recursive_size(obj):
+    """Get the memory referenced from this object.
+
+    This returns the memory of the direct object, and all of the memory
+    referenced by child objects. It also returns the total number of objects.
+    """
+    total_size = 0
+    pending = [obj]
+    last_item = 0
+    seen = _intset.IDSet()
+    size_of = _scanner.size_of
+    while last_item >= 0:
+        item = pending[last_item]
+        last_item -= 1
+        id_item = id(item)
+        if id_item in seen:
+            continue
+        seen.add(id_item)
+        total_size += size_of(item)
+        for child in gc.get_referents(item):
+            if id(child) not in seen:
+                last_item += 1
+                if len(pending) > last_item:
+                    pending[last_item] = child
+                else:
+                    pending.append(child)
+    return len(seen), total_size
+
+
+size_of = _scanner.size_of

=== modified file 'meliae/tests/test_scanner.py'
--- a/meliae/tests/test_scanner.py	2009-09-08 17:19:16 +0000
+++ b/meliae/tests/test_scanner.py	2009-09-08 17:41:03 +0000
@@ -76,3 +76,40 @@
         # We have a reference cycle here, but we should not loop endlessly :)
         self.assertDumpAllReferenced([a, b, c, l], l)
         self.assertDumpAllReferenced([a, b, c, l], c)
+
+
+class TestGetRecursiveSize(tests.TestCase):
+
+    def assertRecursiveSize(self, n_objects, total_size, obj):
+        self.assertEqual((n_objects, total_size),
+                         scanner.get_recursive_size(obj))
+
+    def test_single_object(self):
+        i = 1
+        self.assertRecursiveSize(1, scanner.size_of(i), i)
+        d = {}
+        self.assertRecursiveSize(1, scanner.size_of(d), d)
+        l = []
+        self.assertRecursiveSize(1, scanner.size_of(l), l)
+
+    def test_referenced(self):
+        s1 = 'this is a simple string'
+        s2 = 'and this is another one'
+        s3 = s1 + s2
+        s4 = 'this is a' + ' simple string'# same as s1, but diff object
+        self.assertTrue(s1 is not s4)
+        self.assertTrue(s1 == s4)
+        d = {s1:s2, s3:s4}
+        total_size = (scanner.size_of(s1) + scanner.size_of(s2)
+                      + scanner.size_of(s3) + scanner.size_of(s4)
+                      + scanner.size_of(d))
+        self.assertRecursiveSize(5, total_size, d)
+
+    def test_recursive(self):
+        s1 = 'this is a simple string'
+        s2 = 'and this is another one'
+        l = [s1, s2]
+        l.append(l)
+        total_size = (scanner.size_of(s1) + scanner.size_of(s2)
+                      + scanner.size_of(l))
+        self.assertRecursiveSize(3, total_size, l)