Simple undo decorator for creating an undo chunk around a function within Maya

undo_decorator.py

from functools import wraps
import maya.cmds as cmds


def one_undo(func):
    """
    Decorator - guarantee close chunk.

    type: (function) -> function
    """
    @wraps(func)
    def wrap(*args, **kwargs):
        try:
            cmds.undoInfo(openChunk=True)
            return func(*args, **kwargs)
        except Exception as e:
            raise e
        finally:
            cmds.undoInfo(closeChunk=True)
    return wrap


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


if __name__ == '__main__':
    weird_func()