Rebased dreb's commit to update sizes dictionary

for faster calculation of sizes.
https://github.com/andrejb/leap_mail/commit/8b88e85fab3c2b75da16b16c8d492c001b8076c6

1 files changed, 14 insertions, 1 deletions

diff --git a/src/leap/mail/imap/memorystore.py b/src/leap/mail/imap/memorystore.py
index 195cef7..a99148f 100644
--- a/src/leap/mail/imap/memorystore.py
+++ b/src/leap/mail/imap/memorystore.py
@@ -106,6 +106,12 @@ class MemoryStore(object):
         # Internal Storage: messages
         self._msg_store = {}
 
+        # Sizes
+        """
+        {'mbox, uid': <int>}
+        """
+        self._sizes = {}
+
         # Internal Storage: payload-hash
         """
         {'phash': weakreaf.proxy(dict)}
@@ -347,8 +353,12 @@ class MemoryStore(object):
             for key in seq:
                 if key in store and empty(store.get(key)):
                     store.pop(key)
+
         prune((FDOC, HDOC, CDOCS, DOCS_ID), store)
 
+        # Update memory store size
+        self._sizes[key] = size(self._msg_store[key])
+
     def get_docid_for_fdoc(self, mbox, uid):
         """
         Return Soledad document id for the flags-doc for a given mbox and uid,
@@ -417,6 +427,9 @@ class MemoryStore(object):
             self._new.discard(key)
             self._dirty.discard(key)
             self._msg_store.pop(key, None)
+            if key in self._sizes:
+                del self._sizes[key]
+
         except Exception as exc:
             logger.exception(exc)
 
@@ -958,4 +971,4 @@ class MemoryStore(object):
 
         :rtype: int
         """
-        return size.get_size(self._msg_store)
+        return reduce(lambda x, y: x + y, self._sizes, 0)