Python2.7 import compatibility fix

Author: cmccandless

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):
@@ -154,9 +156,14 @@ def from_EmailMessage(cls, message):
             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(),
-            message.get_content()
+            body.strip()
         ))
         for k, v in message.items():
             mail.add_header(Header(k, v))

test/test_mail.py

@@ -2,7 +2,12 @@
 import json
 import unittest
 
-from email.message import EmailMessage
+try:
+    from email.message import EmailMessage
+except ImportError:
+    # Python2
+    from email import message
+    EmailMessage = message.Message
 
 from sendgrid.helpers.mail import (
     ASM,
@@ -562,7 +567,12 @@ def test_directly_setting_substitutions(self):
 
     def test_from_emailmessage(self):
         message = EmailMessage()
-        message.set_content('message that is not urgent')
+        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'
@@ -576,4 +586,4 @@ def test_from_emailmessage(self):
         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\n')
+        self.assertEqual(content.value, 'message that is not urgent')