Package madgraph :: Package iolibs :: Module save_load_object
[hide private]
[frames] | no frames]

Source Code for Module madgraph.iolibs.save_load_object

  1  ################################################################################ 
  2  # 
  3  # Copyright (c) 2009 The MadGraph5_aMC@NLO Development team and Contributors 
  4  # 
  5  # This file is a part of the MadGraph5_aMC@NLO project, an application which  
  6  # automatically generates Feynman diagrams and matrix elements for arbitrary 
  7  # high-energy processes in the Standard Model and beyond. 
  8  # 
  9  # It is subject to the MadGraph5_aMC@NLO license which should accompany this  
 10  # distribution. 
 11  # 
 12  # For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch 
 13  # 
 14  ################################################################################ 
 15   
 16  """Function to save any Python object to file.""" 
 17   
 18  from __future__ import absolute_import 
 19  import pickle 
 20  import six.moves.cPickle 
 21   
 22  from . import files as files 
 23  import six 
 24  import os 
 25   
26 -class SaveObjectError(Exception):
27 """Exception raised if an error occurs in while trying to save an 28 object to file.""" 29 pass
30
31 -def save_to_file(filename, object, log=True, allow_fail=True):
32 """Save any Python object to file filename""" 33 34 if not isinstance(filename, six.string_types): 35 raise SaveObjectError("filename must be a string") 36 37 files.write_to_file(filename, pickle_object, object, log=log, binary=True, 38 bypass_error=True) 39 40 return True
41
42 -def load_from_file(filename,binary=True):
43 """Save any Python object to file filename""" 44 45 if not isinstance(filename, str): 46 raise SaveObjectError("filename must be a string") 47 return files.read_from_file(filename, unpickle_object, binary=binary)
48
49 -def pickle_object(fsock, object, bypass_error=False, **opts):
50 """Helper routine to pickle an object to file socket fsock""" 51 52 try: 53 six.moves.cPickle.dump(object, fsock, protocol=2) 54 except Exception as error: 55 if bypass_error: 56 return 57 else: 58 raise
59 60
61 -class UnPickler(pickle.Unpickler):
62 """Treat problem of librarie""" 63
64 - def __init__(self, *args, **opts):
65 pickle.Unpickler.__init__(self, *args, **opts) 66 self.basemod = os.path.dirname(args[0].name)
67
68 - def find_class(self, module, name):
69 """Find the correct path for the given function. 70 Due to ME call via MG some libraries might be messed up on the pickle 71 This routine helps to find back which one we need. 72 """ 73 74 # A bit of an ugly hack, but it works and has no side effect. 75 if module == 'loop_me_comparator': 76 module = 'tests.parallel_tests.loop_me_comparator' 77 import sys 78 try: 79 import madgraph.various.misc as misc 80 except ImportError: 81 import internal.misc as misc 82 83 with misc.TMP_variable(sys, 'path', sys.path + [self.basemod]): 84 try: 85 return pickle.Unpickler.find_class(self, module, name) 86 except ImportError as error: 87 pass 88 89 lerror = None 90 for prefix in ['internal.%s', 'madgraph.iolibs.%s', 'madgraph.madevent.%s', 91 'madgraph.various.%s', 'internal.ufomodel.%s']: 92 93 if '.' in module: 94 newmodule = prefix % module.rsplit('.',1)[1] 95 else: 96 newmodule = prefix % module 97 98 try: 99 return pickle.Unpickler.find_class(self, newmodule , name) 100 except Exception as error: 101 lerror = error 102 pass 103 104 else: 105 raise lerror
106 107
108 -def unpickle_object(fsock):
109 """Helper routine to pickle an object to file socket fsock""" 110 111 p = UnPickler(fsock) 112 return p.load()
113 #return pickle.load(fsock) 114