[umessage] Use the same code paths for both python 2 and 3

Add tests for UMessage.get_payload() (taken from CubicWeb).
Closes #1716128.

2 files changed, 49 insertions, 54 deletions

diff --git a/logilab/common/umessage.py b/logilab/common/umessage.py
index 459738d..a0394bc 100644
--- a/logilab/common/umessage.py
+++ b/logilab/common/umessage.py
@@ -101,61 +101,39 @@ class UMessage:
         for part in self.message.walk():
             yield UMessage(part)
 
-    if sys.version_info < (3, 0):
-
-        def get_payload(self, index=None, decode=False):
-            message = self.message
-            if index is None:
-                payload = message.get_payload(index, decode)
-                if isinstance(payload, list):
-                    return [UMessage(msg) for msg in payload]
-                if message.get_content_maintype() != 'text':
-                    return payload
-
-                charset = message.get_content_charset() or 'iso-8859-1'
-                if search_function(charset) is None:
-                    charset = 'iso-8859-1'
-                return unicode(payload or '', charset, "replace")
-            else:
-                payload = UMessage(message.get_payload(index, decode))
-            return payload
-
-        def get_content_maintype(self):
-            return unicode(self.message.get_content_maintype())
-
-        def get_content_type(self):
-            return unicode(self.message.get_content_type())
-
-        def get_filename(self, failobj=None):
-            value = self.message.get_filename(failobj)
-            if value is failobj:
-                return value
-            try:
-                return unicode(value)
-            except UnicodeDecodeError:
-                return u'error decoding filename'
-
-    else:
-
-        def get_payload(self, index=None, decode=False):
-            message = self.message
-            if index is None:
-                payload = message.get_payload(index, decode)
-                if isinstance(payload, list):
-                    return [UMessage(msg) for msg in payload]
+    def get_payload(self, index=None, decode=False):
+        message = self.message
+        if index is None:
+            payload = message.get_payload(index, decode)
+            if isinstance(payload, list):
+                return [UMessage(msg) for msg in payload]
+            if message.get_content_maintype() != 'text':
+                return payload
+            if isinstance(payload, text_type):
                 return payload
-            else:
-                payload = UMessage(message.get_payload(index, decode))
-            return payload
-
-        def get_content_maintype(self):
-            return self.message.get_content_maintype()
-
-        def get_content_type(self):
-            return self.message.get_content_type()
 
-        def get_filename(self, failobj=None):
-            return self.message.get_filename(failobj)
+            charset = message.get_content_charset() or 'iso-8859-1'
+            if search_function(charset) is None:
+                charset = 'iso-8859-1'
+            return text_type(payload or b'', charset, "replace")
+        else:
+            payload = UMessage(message.get_payload(index, decode))
+        return payload
+
+    def get_content_maintype(self):
+        return text_type(self.message.get_content_maintype())
+
+    def get_content_type(self):
+        return text_type(self.message.get_content_type())
+
+    def get_filename(self, failobj=None):
+        value = self.message.get_filename(failobj)
+        if value is failobj:
+            return value
+        try:
+            return text_type(value)
+        except UnicodeDecodeError:
+            return u'error decoding filename'
 
     # other convenience methods ###############################################
 
diff --git a/test/unittest_umessage.py b/test/unittest_umessage.py
index 940b308..2841172 100644
--- a/test/unittest_umessage.py
+++ b/test/unittest_umessage.py
@@ -23,7 +23,7 @@ from os.path import join, dirname, abspath
 from six import text_type
 
 from logilab.common.testlib import TestCase, unittest_main
-from logilab.common.umessage import UMessage, decode_QP
+from logilab.common.umessage import UMessage, decode_QP, message_from_string
 
 DATA = join(dirname(abspath(__file__)), 'data')
 
@@ -53,6 +53,23 @@ class UMessageTC(TestCase):
     def test_get_payload_no_multi(self):
         payload = self.umessage1.get_payload()
         self.assertEqual(type(payload), text_type)
+    
+    def test_get_payload_decode(self):
+        msg = """\
+MIME-Version: 1.0
+Content-Type: text/plain; charset="utf-8"
+Content-Transfer-Encoding: base64
+Subject: =?utf-8?q?b=C3=AFjour?=
+From: =?utf-8?q?oim?= <oim@logilab.fr>
+Reply-to: =?utf-8?q?oim?= <oim@logilab.fr>, =?utf-8?q?BimBam?= <bim@boum.fr>
+X-CW: data
+To: test@logilab.fr
+Date: now
+
+dW4gcGV0aXQgY8O2dWNvdQ==
+"""
+        msg = message_from_string(msg)
+        self.assertEqual(msg.get_payload(decode=True), u'un petit cöucou')
 
     def test_decode_QP(self):
         test_line =  '=??b?UmFwaGHrbA==?= DUPONT<raphael.dupont@societe.fr>'