Capture the local variables of a function upon exception.

capturing_locals.py

from contextlib import contextmanager
import sys


@contextmanager
def capturing_locals(depth=1):
    """
    Capture the local variables of a function upon exception.

    >>> def fun():
    ...     x = 1
    ...     assert 0
    ...     y = 2
    ...
    >>> with capturing_locals() as midway:
    ...     fun()
    ...
    Traceback (most recent call last):
      ...
    AssertionError
    >>> midway
    {'x': 1}

    """
    midway = {}
    try:
        yield midway
    except Exception:
        _type, _value, tb = sys.exc_info()
        for _ in range(depth + 1):  # +1 to skip @contextmanager
            tb = tb.tb_next
        midway.update(tb.tb_frame.f_locals)
        del tb
        raise