2017-08-08

Python in Maya #3: Context Manager


Building on top of Python in Maya #1: Decorators because
from contextlib import contextmanager

import pymel.core as pm


# this is the decorator from the old post
def working_unit_linear_cm(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        scene_unit = pm.currentUnit(query=True, linear=True)
        if scene_unit != 'cm':
            pm.currentUnit(linear='cm')
        try:
            result = func(*args, **kwargs)
        finally:
            if scene_unit != 'cm':
                pm.currentUnit(linear=scene_unit)
        return result
    return wrapper


# context manager version
@contextmanager
def working_unit_linear(value='cm'):
    scene_value = pm.currentUnit(query=True, linear=True)
    if scene_value != value:
        pm.currentUnit(linear=value)
    yield
    if scene_value != value:
        pm.currentUnit(linear=scene_value)


# rewriting the decorator using the new context manager
def working_unit_linear_cm(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        with working_unit_linear():
            result = func(*args, **kwargs)
        return result


# generalized currentUnit context manager
@contextmanager
def current_unit(**flags):
    different_scene_values = {}
    for k, value in flags.iteritems():
        scene_value = pm.currentUnit(query=True, **{k: True})
        if scene_value != value:
            pm.currentUnit(**{k: value})
            different_scene_values[k] = scene_value
    yield
    for k, value in different_scene_values.iteritems():
        pm.currentUnit(**{k: value})


# rewriting the first context manager
@contextmanager
def current_unit_linear(value='cm'):
    with current_unit(linear=value):
        yield
Notes
  • Only use current_unit?
  • Generalize current_unit to any command?
  • Decorator with argument examples?
  • Name working_units or current_unit?
  • Procedural programming paradigm is the "opposite", by extracting/reusing the middle code section (vs. the wrapping part)
  • The ContextDecorator can be used as decorator and context manager: https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator

"""
More contextmanager examples
"""
from contextlib import contextmanager

import pymel.core as pm


@contextmanager
def unlocked_node(node):
    node = pm.PyNode(node)
    is_locked = node.isLocked()
    if is_locked:
        node.setLocked(False)
    yield node
    if is_locked and node.exists():
        node.setLocked(True)


@contextmanager
def unlocked_attribute(attribute):
    attribute = pm.PyNode(attribute)
    with unlocked_node(attribute.node()):
        is_locked = attribute.isLocked()
        if is_locked:
            attribute.setLocked(False)
        yield attribute
        if is_locked and attribute.exists():
            attribute.setLocked(True)


@contextmanager
def attribute_value(attribute, value):
    attribute = pm.PyNode(attribute)
    start_value = attribute.get()
    if start_value == value:
        yield attribute
        return
    with unlocked_attribute(attribute) as attr:
        if start_value != value:
            attr.set(value)
        yield attr
        if start_value != value:
            attr.set(start_value)


# TODO: generalized unlock context manager that detects attr/node?