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

1 
2      def walk(self):
3          for part in self.message.walk():
4              yield UMessage(part)
5 
6 -    if sys.version_info < (3, 0):
7 -
8 -        def get_payload(self, index=None, decode=False):
9 -            message = self.message
10 -            if index is None:
11 -                payload = message.get_payload(index, decode)
12 -                if isinstance(payload, list):
13 -                    return [UMessage(msg) for msg in payload]
14 -                if message.get_content_maintype() != 'text':
15 -                    return payload
16 -
17 -                charset = message.get_content_charset() or 'iso-8859-1'
18 -                if search_function(charset) is None:
19 -                    charset = 'iso-8859-1'
20 -                return unicode(payload or '', charset, "replace")
21 -            else:
22 -                payload = UMessage(message.get_payload(index, decode))
23 -            return payload
24 -
25 -        def get_content_maintype(self):
26 -            return unicode(self.message.get_content_maintype())
27 -
28 -        def get_content_type(self):
29 -            return unicode(self.message.get_content_type())
30 +    def get_payload(self, index=None, decode=False):
31 +        message = self.message
32 +        if index is None:
33 +            payload = message.get_payload(index, decode)
34 +            if isinstance(payload, list):
35 +                return [UMessage(msg) for msg in payload]
36 +            if message.get_content_maintype() != 'text':
37 +                return payload
38 +            if isinstance(payload, text_type):
39 +                return payload
40 
41 -        def get_filename(self, failobj=None):
42 -            value = self.message.get_filename(failobj)
43 -            if value is failobj:
44 -                return value
45 -            try:
46 -                return unicode(value)
47 -            except UnicodeDecodeError:
48 -                return u'error decoding filename'
49 +            charset = message.get_content_charset() or 'iso-8859-1'
50 +            if search_function(charset) is None:
51 +                charset = 'iso-8859-1'
52 +            return text_type(payload or b'', charset, "replace")
53 +        else:
54 +            payload = UMessage(message.get_payload(index, decode))
55 +        return payload
56 
57 -    else:
58 +    def get_content_maintype(self):
59 +        return text_type(self.message.get_content_maintype())
60 
61 -        def get_payload(self, index=None, decode=False):
62 -            message = self.message
63 -            if index is None:
64 -                payload = message.get_payload(index, decode)
65 -                if isinstance(payload, list):
66 -                    return [UMessage(msg) for msg in payload]
67 -                return payload
68 -            else:
69 -                payload = UMessage(message.get_payload(index, decode))
70 -            return payload
71 +    def get_content_type(self):
72 +        return text_type(self.message.get_content_type())
73 
74 -        def get_content_maintype(self):
75 -            return self.message.get_content_maintype()
76 -
77 -        def get_content_type(self):
78 -            return self.message.get_content_type()
79 -
80 -        def get_filename(self, failobj=None):
81 -            return self.message.get_filename(failobj)
82 +    def get_filename(self, failobj=None):
83 +        value = self.message.get_filename(failobj)
84 +        if value is failobj:
85 +            return value
86 +        try:
87 +            return text_type(value)
88 +        except UnicodeDecodeError:
89 +            return u'error decoding filename'
90 
91      # other convenience methods ###############################################
92 
93      def headers(self):
94          """return an unicode string containing all the message's headers"""

95  from os.path import join, dirname, abspath
96 
97  from six import text_type
98 
99  from logilab.common.testlib import TestCase, unittest_main
100 -from logilab.common.umessage import UMessage, decode_QP
101 +from logilab.common.umessage import UMessage, decode_QP, message_from_string
102 
103  DATA = join(dirname(abspath(__file__)), 'data')
104 
105  class UMessageTC(TestCase):
106 

107          self.assertEqual(to, [u'�l�ment � accents <alf@logilab.fr>'])
108 
109      def test_get_payload_no_multi(self):
110          payload = self.umessage1.get_payload()
111          self.assertEqual(type(payload), text_type)
112 +    
113 +    def test_get_payload_decode(self):
114 +        msg = """\
115 +MIME-Version: 1.0
116 +Content-Type: text/plain; charset="utf-8"
117 +Content-Transfer-Encoding: base64
118 +Subject: =?utf-8?q?b=C3=AFjour?=
119 +From: =?utf-8?q?oim?= <oim@logilab.fr>
120 +Reply-to: =?utf-8?q?oim?= <oim@logilab.fr>, =?utf-8?q?BimBam?= <bim@boum.fr>
121 +X-CW: data
122 +To: test@logilab.fr
123 +Date: now
124 +
125 +dW4gcGV0aXQgY8O2dWNvdQ==
126 +"""
127 +        msg = message_from_string(msg)
128 +        self.assertEqual(msg.get_payload(decode=True), u'un petit c�ucou')
129 
130      def test_decode_QP(self):
131          test_line =  '=??b?UmFwaGHrbA==?= DUPONT<raphael.dupont@societe.fr>'
132          test = decode_QP(test_line)
133          self.assertEqual(type(test), text_type)