Update to new (3.6) email message API

This uses email.message.EmailMessage.get_body() to find the best
candidate for a "body" message part and replaces our own ad-hoc solution
in alot.db.utils.extract_bodytext, which was based on a walk through all
parts.

2 files changed, 41 insertions, 58 deletions

diff --git a/alot/db/message.py b/alot/db/message.py
index 76bb9dd0..8b93f1ce 100644
--- a/alot/db/message.py
+++ b/alot/db/message.py
@@ -271,16 +271,20 @@ class Message:
         """
         # TODO: allow toggle commands to decide which part is considered body
         eml = self.get_email()
-        bodytext = extract_body(eml)
-
-        # check if extracted body is empty but msg contains html parts
-        if (not bodytext and
-           'text/html' in (part.get_content_type() for part in eml.walk())):
-            return MISSING_HTML_MSG
+        bodytext = self.get_text_content()
+
+        # check if extracted body is empty dispite having a text/html body part
+        if eml is not None:
+            bodypart = eml.get_body()
+            if (bodypart and
+                    bodypart.get_payload() and
+                    eml.get_content_type() == 'text/html' and
+                    not bodytext):
+                bodytext = MISSING_HTML_MSG
         return bodytext
 
     def get_text_content(self):
-        return extract_body(self.get_email(), types=['text/plain'])
+        return extract_body(self.get_email())
 
     def matches(self, querystring):
         """tests if this messages is in the resultset for `querystring`"""
diff --git a/alot/db/utils.py b/alot/db/utils.py
index 49d5b6d2..da8f3df7 100644
--- a/alot/db/utils.py
+++ b/alot/db/utils.py
@@ -211,7 +211,9 @@ def _handle_encrypted(original, message, session_keys=None):
 
     if malformed:
         msg = u'Malformed OpenPGP message: {0}'.format(malformed)
-        content = email.message_from_string(msg, policy=email.policy.SMTP)
+        content = email.message_from_string(msg,
+                                            _class=email.message.EmailMessage,
+                                            policy=email.policy.SMTP)
         content.set_charset('utf-8')
         original.attach(content)
 
@@ -228,7 +230,8 @@ def decrypted_message_from_file(handle, session_keys=None):
     :returns: :class:`email.message.Message` possibly augmented with
               decrypted data
     '''
-    return decrypted_message_from_message(email.message_from_file(handle),
+    return decrypted_message_from_message(email.message_from_file(handle,
+                                          _class=email.message.EmailMessage),
                                           session_keys)
 
 
@@ -298,7 +301,9 @@ def decrypted_message_from_bytes(bytestring, session_keys=None):
     :param session_keys: a list OpenPGP session keys
     """
     return decrypted_message_from_message(
-        email.message_from_bytes(bytestring, policy=email.policy.SMTP),
+        email.message_from_bytes(bytestring,
+                                 _class=email.message.EmailMessage,
+                                 policy=email.policy.SMTP),
         session_keys)
 
 
@@ -453,63 +458,37 @@ def remove_cte(part, as_string=False):
     return bp
 
 
-def extract_body(mail, types=None, field_key='copiousoutput'):
+def extract_body(mail):
     """Returns a string view of a Message.
 
-    If the `types` argument is set then any encoding types there will be used
-    as the prefered encoding to extract. If `types` is None then
-    :ref:`prefer_plaintext <prefer-plaintext>` will be consulted; if it is True
-    then text/plain parts will be returned, if it is false then text/html will
-    be returned if present or text/plain if there are no text/html parts.
+    This consults :ref:`prefer_plaintext <prefer-plaintext>`
+    to determine if a "text/plain" alternative is preferred over a "text/html"
+    part.
 
     :param mail: the mail to use
     :type mail: :class:`email.Message`
-    :param types: mime content types to use for body string
-    :type types: list[str]
     :returns: The combined text of any parts to be used
     :rtype: str
     """
 
-    preferred = 'text/plain' if settings.get(
-        'prefer_plaintext') else 'text/html'
-    has_preferred = False
-
-    # see if the mail has our preferred type
-    if types is None:
-        has_preferred = list(typed_subpart_iterator(
-            mail, *preferred.split('/')))
-
-    body_parts = []
-    for part in mail.walk():
-        # skip non-leaf nodes in the mail tree
-        if part.is_multipart():
-            continue
-
-        ctype = part.get_content_type()
-
-        if types is not None:
-            if ctype not in types:
-                continue
-        cd = part.get('Content-Disposition', '')
-        if cd.startswith('attachment'):
-            continue
-        # if the mail has our preferred type, we only keep this type
-        # note that if types != None, has_preferred always stays False
-        if has_preferred and ctype != preferred:
-            continue
-
-        if ctype == 'text/plain':
-            body_parts.append(string_sanitize(remove_cte(part, as_string=True)))
-        else:
-            rendered_payload = render_part(part)
-            if rendered_payload:  # handler had output
-                body_parts.append(string_sanitize(rendered_payload))
-            # mark as attachment
-            elif cd:
-                part.replace_header('Content-Disposition', 'attachment; ' + cd)
-            else:
-                part.add_header('Content-Disposition', 'attachment;')
-    return u'\n\n'.join(body_parts)
+    if settings.get('prefer_plaintext'):
+        preferencelist = ('plain',)
+    else:
+        preferencelist = ('html', 'plain')
+
+    body_part = mail.get_body(preferencelist)
+    if body_part is None:  # if no part matching preferredlist was found
+        return ""
+
+    displaystring = ""
+
+    if body_part.get_content_type() == 'text/plain':
+        displaystring = string_sanitize(remove_cte(body_part, as_string=True))
+    else:
+        rendered_payload = render_part(body_part)
+        if rendered_payload:  # handler had output
+            displaystring = string_sanitize(rendered_payload)
+    return displaystring
 
 
 def formataddr(pair):