quick and dirty Maya Python undo decorator

undo_dec.py

from functools import wraps
from maya import cmds

def undo(func):
    """ Puts the wrapped `func` into a single Maya Undo action, then 
        undoes it when the function enters the finally: block """
    @wraps(func)
    def _undofunc(*args, **kwargs):
        try:
            # start an undo chunk
            cmds.undoInfo(ock=True)
            return func(*args, **kwargs)
        finally:
            # after calling the func, end the undo chunk and undo
            cmds.undoInfo(cck=True)
            cmds.undo()

    return _undofunc

usage.py

from undo_dec import undo

from maya import cmds

@undo
def weird_func():
    """ Creates a bunch of spheres and then gets rolled back """
    rad = 0
    for i in xrange(1000):
        print i
        rad += 0.2 * i
        cmds.sphere(r=rad)
    return rad

print weird_func()