Fixes: #7 py2-py3 compat for custom exception

TeamMsgExtractor · Dec 10, 2018 · 4c0ed08 · 4c0ed08
1 parent e2a3fc1
commit 4c0ed08
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/extract_msg/message.py b/extract_msg/message.py
@@ -7,14 +7,15 @@
 
 from imapclient.imapclient import decode_utf7
 import olefile
-from extract_msg.exceptions import InvalidFileFormat
+
 
 from email.parser import Parser as EmailParser
 from extract_msg import constants
 from extract_msg.attachment import Attachment
 from extract_msg.properties import Properties
 from extract_msg.recipient import Recipient
 from extract_msg.utils import addNumToDir, encode, has_len, stri, windowsUnicode, xstr
+from extract_msg.exceptions import InvalidFileFormat
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
@@ -41,11 +42,16 @@ def __init__(self, path, prefix='', attachmentClass=Attachment, filename=None):
         logger.debug('prefix: {}'.format(prefix))
         self.__path = path
         self.__attachmentClass = attachmentClass
+
         try:
             olefile.OleFileIO.__init__(self, path)
-        except OSError as e:
+        except OSError as e:    # py3
             logger.error(e)
             raise InvalidFileFormat
+        except IOError as e:    # py2
+            logger.error(e)
+            raise InvalidFileFormat
+
         prefixl = []
         if prefix != '':
             if not isinstance(prefix, stri):
@@ -58,13 +64,13 @@ def __init__(self, path, prefix='', attachmentClass=Attachment, filename=None):
             g = prefix.split("/")
             if g[-1] == '':
                 g.pop()
-            prefixl = g
+            prefix = g
             if prefix[-1] != '/':
                 prefix += '/'
             filename = self._getStringStream(prefixl[:-1] + ['__substg1.0_3001'], prefix=False)
         self.__prefix = prefix
         self.__prefixList = prefixl
-        if filename != None:
+        if filename is not None:
             self.filename = filename
         elif has_len(path):
             if len(path) < 1536:
@@ -149,7 +155,7 @@ def _getStringStream(self, filename, prefer='unicode', prefix=True):
         asciiVersion = self._getStream(filename + '001E', prefix)
         unicodeVersion = windowsUnicode(self._getStream(filename + '001F', prefix))
         logger.debug('_getStringSteam called for {}. Ascii version found: {}. Unicode version found: {}.'.format(
-            filename, asciiVersion != None, unicodeVersion != None))
+            filename, asciiVersion is not None, unicodeVersion is not None))
         if asciiVersion is None:
             return unicodeVersion
         elif unicodeVersion is None:
@@ -504,7 +510,7 @@ def save(self, toJson=False, useFileName=False, raw=False, ContentId=False, cust
         else:
             if useFileName:
                 # strip out the extension
-                if self.filename != None:
+                if self.filename is not None:
                     dirName = self.filename.split('/').pop().split('.')[0]
                 else:
                     ValueError(