Allow creation of Mail from EmailMessage

sendgrid/helpers/mail/mail.py

@@ -1,6 +1,8 @@
 """v3/mail/send response body builder"""
 from .personalization import Personalization
 from .header import Header
+from .email import Email
+from .content import Content
 
 
 class Mail(object):
@@ -141,3 +143,28 @@ def get(self):
 
         return {key: value for key, value in mail.items()
                 if value is not None and value != [] and value != {}}
+
+    @classmethod
+    def from_EmailMessage(cls, message):
+        """Create a Mail object from an instance of
+        email.message.EmailMessage.
+        :type message: email.message.EmailMessage
+        :rtype: Mail
+        """
+        mail = cls(
+            from_email=Email(message.get('From')),
+            subject=message.get('Subject'),
+            to_email=Email(message.get('To')),
+        )
+        try:
+            body = message.get_content()
+        except AttributeError:
+            # Python2
+            body = message.get_payload()
+        mail.add_content(Content(
+            message.get_content_type(),
+            body.strip()
+        ))
+        for k, v in message.items():
+            mail.add_header(Header(k, v))
+        return mail

test/test_mail.py

@@ -2,6 +2,13 @@
 import json
 import unittest
 
+try:
+    from email.message import EmailMessage
+except ImportError:
+    # Python2
+    from email import message
+    EmailMessage = message.Message
+
 from sendgrid.helpers.mail import (
     ASM,
     APIKeyIncludedException,
@@ -557,3 +564,26 @@ def test_disable_tracking(self):
     def test_directly_setting_substitutions(self):
         personalization = Personalization()
         personalization.substitutions = [{'a': 0}]
+
+    def test_from_emailmessage(self):
+        message = EmailMessage()
+        body = 'message that is not urgent'
+        try:
+            message.set_content(body)
+        except AttributeError:
+            # Python2
+            message.set_payload(body)
+        message.set_default_type('text/plain')
+        message['Subject'] = 'URGENT TITLE'
+        message['From'] = 'test@example.com'
+        message['To'] = 'test@sendgrid.com'
+        mail = Mail.from_EmailMessage(message)
+        self.assertEqual(mail.subject, 'URGENT TITLE')
+        self.assertEqual(mail.from_email.email, 'test@example.com')
+        self.assertEqual(len(mail.personalizations), 1)
+        self.assertEqual(len(mail.personalizations[0].tos), 1)
+        self.assertDictEqual(mail.personalizations[0].tos[0], {'email': 'test@sendgrid.com'})
+        self.assertEqual(len(mail.contents), 1)
+        content = mail.contents[0]
+        self.assertEqual(content.type, 'text/plain')
+        self.assertEqual(content.value, 'message that is not urgent')