CODEWITHSUNDEEP
python
BenDundee
BenDundee

Reputation: 4521

A fail-safe for long-running python scripts

This may be the dumbest question ever, but I'd like to know if there's a way to write a wrapper around a function (preferably a decorator) so that you can catch the internal state of the local variables in the event that an exception was raised in that function. It would catch the locals as they're created, pickle them, and then dispose of them if no exceptions were raised, or write them to file if any exceptions were found.

Is this too fanciful, or has anyone fooled around with something like this?

Upvotes: 3

Views: 1094

Answers (2)

Wesley Baugh
Wesley Baugh

Reputation: 3770

Another option would be to pass as arguments to the function just those variables that you are interested in saving in case of something catastrophic.

A decorator that might help you to do this:

import os
try:
    import cPickle as pickle
except ImportError:
    import pickle

def save_args_if_exception(pickle_folder):
    def real_decorator(function):
        def wrapper(*args, **kwargs):
            try:
                function(*args, **kwargs)
            except:
                print 'FAILING SAFE: Saving arguments to folder:'
                print pickle_folder
                for arg_position, arg in enumerate(args, start=1):
                    arg_fname = os.path.join(pickle_folder, str(arg_position) +
                                             '.pickle')
                    print 'Saving positional argument-{}'.format(arg_position)
                    with open(arg_fname, mode='wb') as f:
                        pickle.dump(arg, f)
                print 'Saving keyword arguments (kwargs)'
                with open(os.path.join(pickle_folder, 'kwargs.pickle'),
                          mode='wb') as f:
                    pickle.dump(kwargs, f)
                raise  # re-raise original exception
        return wrapper
    return real_decorator

@save_args_if_exception(os.getcwd())
def important_function(first_arg, second_arg=2, **kwargs):
    print 'About to raise an exception!'
    raise Exception

important_function(3)

Note that in this example, second_arg is not saved to a file, presumably because of the way Python handles default arguments. Depending on your situation, this may or may not be useful behavior.

References:

Upvotes: 1

Martijn Pieters
Martijn Pieters

Reputation: 1123900

You can capture the f_locals variable on a frame in the traceback:

import sys
import functools

def capturelocals(func):
    @functools.wraps(func)
    def wrapperfunc(*args, **kw):
        try:
            return func(*args, **kw)
        except Exception:
            _, _, tb = sys.exc_info()
            try:
                while tb.tb_next is not None:
                    tb = tb.tb_next  # find innermost frame
                locals = tb.tb_frame.f_locals
                print locals
            finally:
                del tb  # prevent leaking tracebacks
            raise
    return wrapperfunc

To demonstrate that it works:

>>> @capturelocals
... def foobar():
...     foo = 'bar'
...     spam = 'eggs'
...     raise ValueError('Bam!')
... 
>>> foobar()
{'foo': 'bar', 'spam': 'eggs'}
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 5, in wrapperfunc
  File "<stdin>", line 5, in foobar
ValueError: Bam!

Upvotes: 5

Related Questions