Cache email object in Message (#148)

* Cache email object in Message * Email object cache flush in some tests
coddingtonbear · Mar 9, 2018 · e068738 · e068738
1 parent 2b2bdb9
commit e068738
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 13 deletions.
diff --git a/django_mailbox/models.py b/django_mailbox/models.py
@@ -345,6 +345,7 @@ def _get_dehydrated_message(self, msg, record):
 
     def _process_message(self, message):
         msg = Message()
+        msg._email_object = message
         settings = utils.get_settings()
 
         if settings['store_original_message']:
@@ -711,20 +712,22 @@ def get_email_object(self):
            (https://docs.python.org/2/library/email.message.html)
 
         """
-        if self.eml:
-            if self.eml.name.endswith('.gz'):
-                body = gzip.GzipFile(fileobj=self.eml).read()
+        if not hasattr(self, '_email_object'):  # Cache fill
+            if self.eml:
+                if self.eml.name.endswith('.gz'):
+                    body = gzip.GzipFile(fileobj=self.eml).read()
+                else:
+                    self.eml.open()
+                    body = self.eml.file.read()
+                    self.eml.close()
             else:
-                self.eml.open()
-                body = self.eml.file.read()
-                self.eml.close()
-        else:
-            body = self.get_body()
-        if six.PY3:
-            flat = email.message_from_bytes(body)
-        else:
-            flat = email.message_from_string(body)
-        return self._rehydrate(flat)
+                body = self.get_body()
+            if six.PY3:
+                flat = email.message_from_bytes(body)
+            else:
+                flat = email.message_from_string(body)
+            self._email_object = self._rehydrate(flat)
+        return self._email_object
 
     def delete(self, *args, **kwargs):
         """Delete this message and all stored attachments."""

diff --git a/django_mailbox/tests/test_message_flattening.py b/django_mailbox/tests/test_message_flattening.py
@@ -64,6 +64,7 @@ def test_message_handles_rehydration_problems(self):
 
         msg = self.mailbox.process_incoming_message(incoming_email_object)
 
+        del msg._email_object  # Cache flush
         actual_email_object = msg.get_email_object()
 
         self.assertEqual(
@@ -89,6 +90,7 @@ def test_message_content_type_stripping(self):
 
             msg = self.mailbox.process_incoming_message(incoming_email_object)
 
+        del msg._email_object  # Cache flush
         actual_email_object = msg.get_email_object()
 
         self.assertEqual(