Merge remote-tracking branch 'refs/remotes/kali/feature/get-sizes' into develop

1 files changed, 14 insertions, 1 deletions

diff --git a/src/leap/mail/imap/memorystore.py b/src/leap/mail/imap/memorystore.py
index 3eba59a..ed2b3f2 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)}
@@ -348,8 +354,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,
@@ -418,6 +428,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)
 
@@ -981,4 +994,4 @@ class MemoryStore(object):
 
         :rtype: int
         """
-        return size.get_size(self._msg_store)
+        return reduce(lambda x, y: x + y, self._sizes, 0)