Merge pull request #2626 from anibalsolon/enh/crash_file_versioning

ENH: Add versioning metadata to crash files

Author: mgxd

nipype/pipeline/plugins/tools.py

@@ -58,7 +58,8 @@ def report_crash(node, traceback=None, hostname=None):
     if crashfile.endswith('.txt'):
         crash2txt(crashfile, dict(node=node, traceback=traceback))
     else:
-        savepkl(crashfile, dict(node=node, traceback=traceback))
+        savepkl(crashfile, dict(node=node, traceback=traceback),
+                versioning=True)
     return crashfile
 
 

nipype/utils/filemanip.py

@@ -628,13 +628,13 @@ def load_json(filename):
 
 
 def loadcrash(infile, *args):
-    if '.pkl' in infile:
-        return loadpkl(infile)
+    if infile.endswith('pkl') or infile.endswith('pklz'):
+        return loadpkl(infile, versioning=True)
     else:
         raise ValueError('Only pickled crashfiles are supported')
 
 
-def loadpkl(infile):
+def loadpkl(infile, versioning=False):
     """Load a zipped or plain cPickled file
     """
     fmlogger.debug('Loading pkl: %s', infile)
@@ -643,11 +643,44 @@ def loadpkl(infile):
     else:
         pkl_file = open(infile, 'rb')
 
+    if versioning:
+        pkl_metadata = {}
+
+        # Look if pkl file contains version file
+        try:
+            pkl_metadata_line = pkl_file.readline()
+            pkl_metadata = json.loads(pkl_metadata_line)
+        except:
+            # Could not get version info
+            pkl_file.seek(0)
+
     try:
-        unpkl = pickle.load(pkl_file)
-    except UnicodeDecodeError:
-        unpkl = pickle.load(pkl_file, fix_imports=True, encoding='utf-8')
-    return unpkl
+        try:
+            unpkl = pickle.load(pkl_file)
+        except UnicodeDecodeError:
+            unpkl = pickle.load(pkl_file, fix_imports=True, encoding='utf-8')
+
+        return unpkl
+
+    # Unpickling problems    
+    except Exception as e:
+        if not versioning:
+            raise e
+
+        from nipype import __version__ as version
+
+        if 'version' in pkl_metadata:
+            if pkl_metadata['version'] != version:
+                fmlogger.error('Your Nipype version is: %s',
+                               version)
+                fmlogger.error('Nipype version of the pkl is: %s',
+                               pkl_metadata['version'])
+        else:
+            fmlogger.error('No metadata was found in the pkl file.')
+            fmlogger.error('Make sure that you are using the same Nipype'
+                           'version from the generated pkl.')
+
+        raise e
 
 
 def crash2txt(filename, record):
@@ -682,11 +715,19 @@ def read_stream(stream, logger=None, encoding=None):
     return out.splitlines()
 
 
-def savepkl(filename, record):
+def savepkl(filename, record, versioning=False):
     if filename.endswith('pklz'):
         pkl_file = gzip.open(filename, 'wb')
     else:
         pkl_file = open(filename, 'wb')
+
+    if versioning:
+        from nipype import __version__ as version
+        metadata = json.dumps({'version': version})
+
+        pkl_file.write(metadata.encode('utf-8'))
+        pkl_file.write('\n'.encode('utf-8'))
+
     pickle.dump(record, pkl_file)
     pkl_file.close()
 

nipype/utils/tests/test_filemanip.py

@@ -8,13 +8,15 @@
 import time
 import warnings
 
+import mock
 import pytest
 from ...testing import TempFATFS
 from ...utils.filemanip import (
     save_json, load_json, fname_presuffix, fnames_presuffix, hash_rename,
     check_forhash, _parse_mount_table, _cifs_table, on_cifs, copyfile,
     copyfiles, ensure_list, simplify_list, check_depends,
-    split_filename, get_related_files, indirectory)
+    split_filename, get_related_files, indirectory,
+    loadpkl, loadcrash, savepkl)
 
 
 def _ignore_atime(stat):
@@ -521,3 +523,50 @@ def test_indirectory(tmpdir):
     except ValueError:
         pass
     assert os.getcwd() == tmpdir.strpath
+
+
+def test_pklization(tmpdir):
+    tmpdir.chdir()
+
+    exc = Exception("There is something wrong here")
+    savepkl('./except.pkz', exc)
+    newexc = loadpkl('./except.pkz')
+
+    assert exc.args == newexc.args
+    assert os.getcwd() == tmpdir.strpath
+
+
+class Pickled:
+
+    def __getstate__(self):
+        return self.__dict__
+
+
+class PickledBreaker:
+
+    def __setstate__(self, d):
+        raise Exception()
+
+        
+def test_versioned_pklization(tmpdir):
+    tmpdir.chdir()
+
+    obj = Pickled()
+    savepkl('./pickled.pkz', obj, versioning=True)
+
+    with pytest.raises(Exception):
+        with mock.patch('nipype.utils.tests.test_filemanip.Pickled', PickledBreaker), \
+             mock.patch('nipype.__version__', '0.0.0'):
+
+            loadpkl('./pickled.pkz', versioning=True)
+
+        
+def test_unversioned_pklization(tmpdir):
+    tmpdir.chdir()
+
+    obj = Pickled()
+    savepkl('./pickled.pkz', obj)
+
+    with pytest.raises(Exception):
+        with mock.patch('nipype.utils.tests.test_filemanip.Pickled', PickledBreaker):
+            loadpkl('./pickled.pkz', versioning=True)