How can I author an undo context manager?

Starting with Python 2.5, they introduced the new with statement, which is known as a 'context manager'. See official docs:

• http://docs.python.org/reference/compound_stmts.html#with
• http://docs.python.org/reference/datamodel.html#context-managers
• http://docs.python.org/library/stdtypes.html#context-manager-types
• http://www.python.org/dev/peps/pep-0343/

And notes on my Python Wiki:

• http://pythonwiki.tiddlyspot.com/#%5B%5Bwith%20statement%5D%5D

In a nutshell, it lets you wrap a block of code in a 'context' that controls how that block will be entered, and exited. It also catches exceptions for you, and returns any that occurred, so it takes care of the try\except clause automatically.

I realized this would be a great way to wrap a block of code in an undo queue, which is often needed when authoring Python\API code in Maya. If this isn't done (based on certain circumstances, usually commands attached to button-press callbacks) when the user undo's, they have to undo every command executed based on the undone function. And if the function had a hundred loops, the user has to press undo a hundred times... Wrapping the code in a undoInfo open\close chunk call, it solves for this. In addition, you need to put your undoInfo calls inside a try\except clause to guarantee that no matter what, you'll close the undo chunk (or your users could get pretty upset with you), and as mentioned above, the with statement takes care of this for you.

I have a solution to this via Python decorators here:
How can I author an undo decorator?
But why not have yet another? :-)

import maya.cmds as mc
import maya.OpenMaya as om

class UndoManager(object):
    """
    This is the context manager class.
    """
    def __enter__(self):
        # Open our undo chunk
        mc.undoInfo(openChunk=True)

    def __exit__(self, exc_type, exc_value, traceback):
        # Close the undo chunk, warn if any exceptions were caught:
        mc.undoInfo(closeChunk=True)
        if exc_type:
            om.MGlobal.displayWarning('%s : %s'%(exc_type, exc_value))
        # If this was false, it would re-raise the exception when complete
        return True 

Simple example usage:

def myFunc():
    print "Executing some code... then:"
    with UndoManager():
        print "Running code in the UndoManager!"
        raise Exception("Oops, an exception :-(")
    print "Finish executing code..."

myFunc()

When executed prints:

Executing some code... then:
Running code in the UndoManager!
# Warning: <type 'exceptions.Exception'> : Oops, an exception :-( # 
Finish executing code...

With UI's:

This works really well when authoring UI's as well. In the below example, we create a UI with button that will create a sphere for every node in the scene, wrapping the sphere creation in our UndoManager():

import maya.cmds as mc

class App(object):

    def __init__(self):
        if mc.window('ucmt', exists=True):
            mc.deleteUI('ucmt')
        mc.window("ucmt", title="Undo ContextManager Test",  resizeToFitChildren=True)
        mc.columnLayout(adjustableColumn=True, columnAttach=('both', 5))
        mc.button(label='FOO!', command=self.buttonCmd)
        mc.showWindow()

    def buttonCmd(self, *args):
        # Wrap everything below in our context manager:
        with UndoManager():
            for i in range(len(mc.ls())):
                mc.polySphere()

# launch the window
App()

If you press the button to make the spheres, you can undo the operation with a single call to undo. However, if you modify the above code and comment out this line:

        with UndoManager():

and try again, you'll see that you'll have to no undo for every sphere that was made. Yuck!