generalise mx/datetime switch to be used by umessage

1 files changed, 15 insertions, 9 deletions

diff --git a/umessage.py b/umessage.py
index 4aa52fe..7766bdd 100644
--- a/umessage.py
+++ b/umessage.py
@@ -10,11 +10,15 @@ import email
 from encodings import search_function
 from email.Utils import parseaddr, parsedate
 from email.Header import decode_header
+from datetime import datetime
 
 try:
     from mx.DateTime import DateTime
 except ImportError:
-    def DateTime(*args): return None
+    DateTime = datetime
+
+import logilab.common as lgc
+
 
 def decode_QP(string):
     parts = []
@@ -30,13 +34,13 @@ def message_from_file(fd):
         return UMessage(email.message_from_file(fd))
     except email.Errors.MessageParseError:
         return ''
-    
+
 def message_from_string(string):
     try:
         return UMessage(email.message_from_string(string))
     except email.Errors.MessageParseError:
         return ''
-    
+
 class UMessage:
     """Encapsulates an email.Message instance and returns only unicode objects.
     """
@@ -45,7 +49,7 @@ class UMessage:
         self.message = message
 
     # email.Message interface #################################################
-    
+
     def get(self, header, default=None):
         value = self.message.get(header, default)
         if value:
@@ -55,7 +59,7 @@ class UMessage:
     def get_all(self, header, default=()):
         return [decode_QP(val) for val in self.message.get_all(header, default)
                 if val is not None]
-    
+
     def get_payload(self, index=None, decode=False):
         message = self.message
         if index is None:
@@ -82,7 +86,7 @@ class UMessage:
     def walk(self):
         for part in self.message.walk():
             yield UMessage(part)
-    
+
     def get_content_maintype(self):
         return unicode(self.message.get_content_maintype())
 
@@ -116,9 +120,9 @@ class UMessage:
             name, mail = parseaddr(person)
             persons.append((name, mail))
         return persons
-    
+
     def date(self, alternative_source=False, return_str=False):
-        """return a mx.DateTime object for the email's date or None if no date is
+        """return a datetime object for the email's date or None if no date is
         set or if it can't be parsed
         """
         value = self.get('date')
@@ -132,7 +136,9 @@ class UMessage:
         if value is not None:
             datetuple = parsedate(value)
             if datetuple:
-                return DateTime(*datetuple[:6])
+                if lgc.USE_MX_DATETIME:
+                    return DateTime(*datetuple[:6])
+                return datetime(*datetuple[:6])
             elif not return_str:
                 return None
         return value