How to group actions into a single Undo Chunk in Maya?

When working with Maya scripts, there are times when we want to group multiple operations into a single undoable action. This allows the artist to undo a series of related changes in one step, providing a more intuitive and efficient workflow. Maya offers the cmds.undoInfo command for managing undo chunks, but improper use of the openChunk and closeChunk flags can leave the undo queue in a bad state which might result in unexpected behaviors or crashes. Therefore, we always want to ensure that we properly close an undo chunk after opening it.

One way to accomplish this, is to wrap our code within a try/except/finally. This will call closeChunk even if an exception is raised:

try:
    cmds.undoInfo(openChunk=True, infinity=True, chunkName='MyUndoChunk')
    # Rest of your code here.
except Exception as e:
    print(e)
finally:
    cmds.undoInfo(closeChunk=True)

Nonetheless, there is an easier way, namely if we use a context manager. You've probably seen this pattern before, where you use a with statement to manage resources that need to be cleaned up, such as file handlers (e.g., with open('file.txt', 'r') as f:). Another option is to use a decorator, which allows you to wrap an entire function with additional functionality without modifying the function itself. Now, Python provides a way to combine these two patterns using the ContextDecorator class, offering flexibility and reusability in our code:

from maya import cmds
from contextlib import ContextDecorator

class Undo(ContextDecorator):
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        cmds.undoInfo(openChunk=True, infinity=True, chunkName=self.name)

    def __exit__(self, exc_type, exc_value, traceback):
        cmds.undoInfo(closeChunk=True)

To use our Undo ContextDecorator as a context manager, we can wrap our code within a with statement. This is useful if we want to wrap a certain block of code within an undo chunk:

from maya import cmds
import random

def run_simulation():
    with Undo("MoveCube"):
        for _ in range(100):
            random_position = random.choices(range(10), k=3)
            cmds.move(*random_position, "pCube1")

If you want to apply the Undo ContextDecorator to an entire function, you can simply use it as a decorator instead:

from maya import cmds
import random

@Undo("MoveCube")
def run_simulation():
    for _ in range(100):
    random_position = random.choices(range(10), k=3)
    cmds.move(*random_position, "pCube1")

A simple demonstration: let's imagine we have a function that moves a cube to a random position 10 consecutive times in the scene. If the artist wants to undo this operation, they would have to undo it 10 times to revert each move. What if our operation involves moving the cube 100 times? Therefore, we want to group this into a single undo chunk so that the artist can undo the move in one step as seen in here:

Happy Coding!