interactive mayapy subprocess

MayaPyHost

"""
MayaPy.py

Remote control of a mayapy session with a simple proxy interface.  This is not intended as a complete RPC server or anything of the sort; the primary use case is as an isolation chamber (since you can manipulate paths and environment variables) for use in testing.

"""

import subprocess
import cPickle as pickle
import os
import sys
# this script is inject to the mayapy at startup time to support the
# interactive repl.
MAYA_INIT = '''
# interactive stubs for controlling a MayaPy instance
# interactively in a subprocess

import atexit
import sys
import traceback
import cPickle as pickle
import importlib

import logging
import maya.standalone
maya.standalone.initialize()
atexit.register(maya.standalone.uninitialize)
import maya.cmds as cmds



def reply(success, value):

    try:
        if hasattr(value, 'decode'):
            value = value.decode('utf-8')
        pickle.dump((success, value), sys.__stdout__, protocol=2)

    finally:
        sys.__stdout__.flush()

def m_eval(cmd):
    try:
        result = eval(cmd)
        reply(True, result or '')

    except Exception as e:
        tb = traceback.format_exc()
        reply (False, tb)



def m_exec(cmd):
    cmd_name = cmd.partition("(")[0].split()[-1] + "()"
    try:
        bc = compile(cmd, "<string>", "exec")
        exec bc in globals()
        reply (True, cmd_name)

    except Exception as e:
        tb = traceback.format_exc()
        reply(False, tb)



def m_proxy(package):
    reply(True, package)
    cmdstring, arg, kwargs = pickle.loads(package)
    cmd = globals()[cmdstring]
    try:
        result = cmd(*arg, **kwargs)
        reply(True, result or '')

    except Exception as e:
        tb = traceback.format_exc()
        reply (False, tb)


def m_get(address):
    value = globals().get(address, None)
    if value is None:
        reply( False, None)
    else:
        reply (True, value)


def m_import(module):
    try: 
        globals()[module] = importlib.import_module(module)
        reply (True, module)
    except Exception as e:
        tb = traceback.format_exc()
        reply (False, tb)

'''


def format_cmd(cmd, *args, **kwargs):
    """
    format a command into an evaluable string.  For some reason Maya seems to care
    about single quotes on the receiving end, so be careful if you touch the 
    'stringify' function.
    """
    result = cmd + '('

    def stringify(x):
        if isinstance(x, str):
            return "'" + x + "'"
        else:
            return str(x)
    args = ", ".join((stringify(a) for a in args))
    kwargs = ", ".join((k + " = " + stringify(v) for k, v in kwargs.items()))
    if args:
        result += args
        if kwargs:
            result += ", "
    if kwargs:
        result += kwargs
    result += ")"

    return result


class ProxyCommandError(RuntimeError):
    """
    When a proxy command is called and fails on the remote host, raise this on the local client

    The exception data will contain the original error
    """
    pass


class CommandProxy(object):
    """
    Syntax sugar for formatting commands on the local client
    """

    def __init__(self, cmd, mayapy):
        self.mayapy = mayapy
        self.cmd = cmd

    def __call__(self, *args, **kwargs):
        cmdstring = format_cmd(self.cmd, *args, **kwargs)
        result, err = self.mayapy.evaluate(cmdstring)
        if err is None:
            return result
        raise ProxyCommandError(err)

    def __getattr__(self, name):
        return CommandProxy(self.cmd + "." + name, self.mayapy)

    def value(self):
        """ if this proxy is pointing at a name, rather than a function, call 'get()' and return its value"""
        val, err = self.mayapy.get(self.cmd)
        if err is None:
            return val
        raise ValueError("Unable to retrieve value")


class MayaPyHost(object):
    """
    Context manager that opens a remote MayaPy instance controls it via
    pipes.

        with MayaPyHost("C:/Program Files/Autodesk/Maya2019/bin/mayapy.exe") as remote:
            remote.evaluate("cmds.polyCube()")

    The primary tools for controlling the remote are

    * evaluate(): calls `eval` on a single expression on  the remote, returning the result
    * execute(): calls `exec` on the remote, so can be used to define a function. Returns true if the execution worked, false if there was an exception on the host.
    * get() :  given a fully qualified name, returns the value of that name on the host

    The MayaPyHost can also create proxy commands to simplify sending commands.  When the context manager
    is active, you can use dot property access to get a proxy object representing a function on the 
    host.  For example:

        with MayaPyHost("C:/Program Files/Autodesk/Maya2019/bin/mayapy.exe") as remote:
            ls = remote.cmds.ls
            print ls()

    is equivalent to 

            remote.evaluate('cmds.ls()')

    The proxy form is usually easier to work with because it is less prone to string formatting and typo problems. Proxies return a value o

    """

    def __init__(self, exe, *paths, **overrides):
        self.exe = exe
        self.history = []
        self.input = None
        self.output = None

        # the supplied paths will replace the PYTHONPATH
        # in the host.  PATH will be emptied out. This may
        # have issues for modules which rely on PATH; you can
        # selectively add path variables back in by including
        # a PATH in your overrides
        runtime_env = os.environ.copy()
        runtime_env['MAYA_LOCATION'] = os.path.dirname(exe)
        runtime_env['PYTHONHOME'] = os.path.dirname(exe)

        def quoted(x):
            return '%s' % os.path.normpath(x)
        # use PYTHONPATH in preference to PATH
        runtime_env['PYTHONPATH'] = ";".join((quoted(p) for p in paths))
        runtime_env['PATH'] = ''
        runtime_env.update(overrides)
        self.environ = dict([(str(k), str(v)) for k, v in runtime_env.items()])

    def __enter__(self):
        # for future reference the '-u' flag forces the pipes into binary mode,
        # which is faster. If you drop -u you'll need to downgrade the protocol
        # in the pickler to protocol 0
        args = self.exe, "-u", "-i", "-c", MAYA_INIT
        self.session = subprocess.Popen(args,
                                        env=self.environ,
                                        bufsize=1,
                                        stdin=subprocess.PIPE,
                                        stdout=subprocess.PIPE, 
                                        stderr =subprocess.PIPE
                                        )

        self.input = self.session.stdin
        self.output = self.session.stdout

        return self

    def communicate(self, cmdstring):
        """
        send cmdstring to the host, and collect the result -- a tuple of value, error
        where value is the result of a valid call and error is the string traceback
        in case of an exception.  These are added to the history of this remote.

        Generally you don't want to call this directly, it's called in either
        `evaluate()` or `execute()` or via a proxy
        """
        self.input.write(cmdstring)
        self.input.flush()
        try:
            result = pickle.load(self.output)
        except pickle.UnpicklingError as e:
            result = (False, e)

        if result[0]:
            self.history.append((result[1], None))
        else:
            self.history.append((None, result[1]))

        return self.history[-1]

    def evaluate(self, cmd):
        """
        evaluate the string 'cmd' on the remote host and return its result. returns
        a tuple value, error where where error is the string traceback in the case of an
        exception
        """
        cmdstring = 'm_eval(r"""{}""")\n'.format(cmd)
        return self.communicate(cmdstring)

    def get(self, addr):
        """
        get the value of the name 'addr' on the remote.  Returns a tuple value, error

        """
        cmdstring = 'm_get(r"""{}""")\n'.format(addr)
        return self.communicate(cmdstring)

    def remote_import(self, module):
        cmdstring = 'm_import(r"""{}""")\n'.format(module)
        try:
            result = self.communicate(cmdstring)
            if result[0]:
                return result[0]
            raise ImportError(result[1])
        except pickle.UnpicklingError:
            raise ImportError("child import failed")

    def execute(self, cmd):
        """
        use exec() on the remote instance.  This allows you to execute functions or
        assign to variables.

        returns true if the code executed, false if it excepted on the host
        """
        cmdstring = 'm_exec(r"""{}""")\n'.format(cmd)
        _, err = self.communicate(cmdstring)
        return err is None

    def __getattr__(self, name):
        """
        We try to provide a command proxy for any dotted property name which is not
        part of the class, eg:

        with MayaPyHost("/path/to/mayapy.exe") as remote:
            ls = remote.cmds.ls
            cube = remote.cmds.polyCube

        the host instance would correspond to the global namespace on the host, so
        'remote.fred' would be a variable named 'fred' in the global namespace
        """

        return CommandProxy(name, self)

    def __exit__(self, *args):
        # close the context manager and shut down the remote mayapy
        self.input.write('sys.exit(0)')
        self.input.flush()
        return False


if __name__ == '__main__':

    # limited self testing
    with MayaPyHost("C:/Program Files/Autodesk/Maya2019/bin/mayapy.exe") as remote:

        # two different ways to issue commands.  The string form...
        result, err = remote.evaluate('cmds.polySphere()')

        # and useiing a proxy
        result = remote.cmds.polyCube()
        assert result == [u'pCube1', u'polyCube1']

        # validate they both worked
        shapes = remote.cmds.ls(type='mesh')
        assert len(shapes) == 2

        # define some functions on the host side
        test_def = """
def test1():
    return 1234

def test2(arg):
    return [1234] * arg
"""
        remote.execute(test_def)

        # make sure they work
        assert remote.evaluate('test1()') == (1234, None)
        assert remote.evaluate('test2(4)') == ([1234, 1234, 1234, 1234], None)

        # setting value on remote requires execute
        remote.execute('fred = 999')
        # retrieve with 'get()'
        assert remote.get('fred') == (999, None)
        remote.execute('fred = 888')
        # or retrieve with a proxy and its' value() method
        assert remote.fred.value() == 888