When working with Maya scripts, you might need to track newly created nodes during specific operations, such as when creating a temporary setup that you want to clean up afterward. Manually tracking nodes by adding them one by one to a list can be cumbersome and error-prone, especially with complex setups. To simplify this process, you can use Maya's addNodeAddedCallback wrapped in a context manager:

import maya.OpenMaya as om

class CaptureNodes:
    def __init__(self, callback, node_type="dependNode"):
        self.callback = callback
        self.node_type = node_type

    def __enter__(self):
        self.id = om.MDGMessage.addNodeAddedCallback(self._callback, self.node_type)

    def _callback(self, node, *_):
        node_name = om.MFnDependencyNode(node).name()
        self.callback(node_name)

    def __exit__(self, *args):
        om.MMessage.removeCallback(self.id)

The context manager handles the setup and teardown of the callback, making your main code cleaner and more focused on the actual node creation. It ensures that your callback is registered when you enter a block of code and automatically removed when you exit it, regardless of whether an exception was raised. This approach is reusable and flexible, allowing you to specify the type of nodes you want to capture and the action you want to perform on them. Here’s an example on how to use it:

import maya.cmds as cmds

my_nodes = []
callback = lambda node: my_nodes.append(node)

with CaptureNodes(callback):
    cmds.createNode("transform")
    cmds.createNode("multiplyDivide")
    cmds.createNode("joint")

print(my_nodes)  # ['transform1', 'multiplyDivide1', 'joint1']

In this example, you create a list my_nodes to store the names of newly created nodes. You define a simple lambda function as your callback, which appends each new node name to the list. The CaptureNodes context manager wraps your node creation operations. After the with block, my_nodes contains the names of all nodes created within the block.