cPickle - sharing pickled objects between scripts and imports

[Rotwang]

Hi all, I have a module that saves and loads data using cPickle, and 
I've encountered a problem. Sometimes I want to import the module and 
use it in the interactive Python interpreter, whereas sometimes I want 
to run it as a script. But objects that have been pickled by running the 
module as a script can't be correctly unpickled by the imported module 
and vice-versa, since how they get pickled depends on whether the 
module's __name__ is '__main__' or 'mymodule' (say). I've tried to get 
around this by adding the following to the module, before any calls to 
cPickle.load:

if __name__ == '__main__':
     import __main__
     def load(f):
         p = cPickle.Unpickler(f)
         def fg(m, c):
             if m == 'mymodule':
                 return getattr(__main__, c)
             else:
                 m = __import__(m, fromlist = [c])
                 return getattr(m, c)
         p.find_global = fg
         return p.load()
else:
     def load(f):
         p = cPickle.Unpickler(f)
         def fg(m, c):
             if m == '__main__':
                 return globals()[c]
             else:
                 m = __import__(m, fromlist = [c])
                 return getattr(m, c)
         p.find_global = fg
         return p.load()
cPickle.load = load
del load


It seems to work as far as I can tell, but I'll be grateful if anyone 
knows of any circumstances where it would fail, or can suggest something 
less hacky. Also, do cPickle.Pickler instances have some attribute 
corresponding to find_global that lets one determine how instances get 
pickled? I couldn't find anything about this in the docs.


-- 
Hate music? Then you'll hate this:

http://tinyurl.com/psymix