Merge a4a3270 into ab8edf2

olefile/olefile.py

@@ -1022,6 +1022,7 @@ def __init__(self, filename=None, raise_defects=DEFECT_FATAL,
             Unicode by default on Python 3+, UTF-8 on Python 2.x.
             (new in olefile 0.42, was hardcoded to Latin-1 until olefile v0.41)
         """
+        self.filename = filename
         # minimal level for defects to be raised as exceptions:
         self._raise_defects_level = raise_defects
         #: list of defects/issues not raised as exceptions:
@@ -1071,6 +1072,15 @@ def __init__(self, filename=None, raise_defects=DEFECT_FATAL,
             self.open(filename, write_mode=write_mode)
 
 
+    def __enter__(self):
+        self.open(self.filename, write_mode=self.write_mode)
+        return self
+
+
+    def __exit__(self, *args):
+        self.close()
+
+
     def _raise_defect(self, defect_level, message, exception_type=IOError):
         """
         This method should be called for any defect found during file parsing.

tests/test_olefile.py

@@ -20,6 +20,11 @@ def test_isOleFile_true(self):
         is_ole = olefile.isOleFile(self.ole_file)
         self.assertTrue(is_ole)
 
+    def test_context_manager(self):
+        with olefile.OleFileIO(self.ole_file) as ole:
+            exists = ole.exists('worddocument')
+            self.assertTrue(exists)
+
     def test_exists_worddocument(self):
         ole = olefile.OleFileIO(self.ole_file)
         exists = ole.exists('worddocument')