lab/m2b/m2b.py

import sys
import os
import pymel.core as pm
import maya.cmds as cmds
import tempfile
import json
import subprocess

# // SETTINGS
GEO_EXT = '.obj'
APP_EXT = '.blend'
PORT = 6006
EDITOR = '/Applications/blender.app/Contents/MacOS/blender'


# shader-independent way of figuring out texture file from mat

# /// UTILS

def guess_diffuse_file(node):
    mat = [pm.ls(pm.listConnections(se), materials=True)[0] for se in node.getShape().outputs(type='shadingEngine')]
    textures = [c for c in mat[0].listConnections() if c.type() =='file']
    for t in textures:
        plugs = t.listConnections(plugs=True)
        if any ([p.name() for p in plugs if 'color' in p.name() or 'diff' in p.name()]):
            return t.fileTextureName.get()
    return 'untextured'


def get_blender_path(root='/Applications'):
    bin_location = 'Contents/MacOS/blender'
    for directory in os.listdir(root):
        full_path = os.path.join(root, directory)
        if os.path.isdir(full_path):
            if directory.endswith('blender.app'):
                return os.path.join(root, directory, bin_location)
            else:
                for subdir in os.listdir(full_path):
                    if subdir.endswith('blender.app'):
                        return os.path.join(full_path, subdir, bin_location)


def this_script():
    # does not work - can we get the current script's location?
    return os.path.join(sys.argv[0])

def replace_mesh(old, new):
    shading_engines = old.getShape().outputs(type='shadingEngine')[0]

    old_mesh = old.listRelatives(shapes=True)[0]
    old_mesh_name = old_mesh.name()
    pm.delete(old_mesh)
    print new
    new_mesh = new.listRelatives(shapes=True)[0]
    pm.rename(new_mesh, old_mesh_name)
    pm.parent(new_mesh, old, shape=True, relative=True)
    pm.delete(new)
    
    pm.sets(shading_engines, edit=True, forceElement=old)  


def load_geofile(filepath):
    try:
        imported_nodes = pm.importFile(filepath, i=True, returnNewNodes=True)
        # print imported_nodes
    
        relevant_nodes = []
        trash = []
        for node in imported_nodes:
            print node, 'TYPE', node.type()
            if node.type() == 'mesh' or node.type() == 'transform':
                relevant_nodes.append(node)
                print '> append', node
            else:
                trash.append(str(node.name()))
                
        # seems like a bug in pymel, materials et al must be
        # explicitly deleted by string name with cmds
        for item in trash:
            try:
                cmds.delete(item)
            except ValueError:
                pass 
            
        for node in relevant_nodes:
            if node.getParent() is None:
                return node 
                
    except RuntimeError:
        print "> unreadable:", filepath
        return False



def open_port(number):
    if not cmds.commandPort(':'+str(number), q=True):
        cmds.commandPort('mayaCommand', name=':'+str(number), sourceType='python')

def get_temp_dir():
    prefix = 'm2b'
    base_tmp = tempfile.gettempdir()
    m2b_tmp = os.path.join(base_tmp, prefix)
    if not os.path.isdir(m2b_tmp):
        os.makedirs(m2b_tmp)
    return m2b_tmp
    
    
# /// CLASSES

class AtOrigin(object):
    def __init__(self, node):
        self.node = node
        self.matrix = pm.xform(matrix=True, query=True)
    
    def __enter__(self):
        print '> discard trs'
        mat_neutral = [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]
        pm.xform(self.node, matrix=mat_neutral)

    def __exit__(self, type, value, traceback):
        print '> restore trs'
        pm.xform(self.node, matrix=self.matrix)


class KeepSelection(object):
    def __enter__(self):
        print '> store sel'
        self.selection = pm.selected()
        
    def __exit__(self, type, value, traceback):
        print '> restore sel'
        pm.select(self.selection)


class BridgedNode(object):
    def __init__(self, uuid):
        self.uuid = uuid
        self.node = None
        self.texture = None
        self.geofile = None
        self.metadatafile = None
        self.matrix = None
        self.valid = False
        self.blendfile = False
        self.updatescript = None   
        
        self.ingest()
    
    def is_valid(self):
        print '> validating:'
        #print [ (key, self.__dict__[key]) for key in self.__dict__.keys()]
    
        if not any([self.__dict__[key] is None for key in self.__dict__.keys()]):
            print '> valid'
            return True
        else:
            print '> invalid'
            print [self.__dict__[key] for key in self.__dict__.keys()]
        
    def has_history(self):
        if os.path.isfile(self.geofile) and os.path.isfile(self.metadatafile) and os.path.isfile(self.blendfile):            
            return True 
    
    def ingest(self):
        # get scene object
        for node in pm.ls(transforms=True):
            if cmds.ls(node.name(), uuid=True)[0] == self.uuid:
                self.node = node
        self.geofile = os.path.join(get_temp_dir(), self.uuid + GEO_EXT)
        self.blendfile = os.path.join(get_temp_dir(), self.uuid + APP_EXT)
        self.metadatafile = os.path.join(get_temp_dir(), self.uuid + '.json')
        self.matrix = pm.xform(matrix=True, query=True)
        self.texture = guess_diffuse_file(self.node)
        self.updatescript = this_script()
        
    def dump_metadata(self):
        metadata = {'obj': self.geofile,
            'tex': self.texture,
            'blend': self.blendfile,
            'uuid': self.uuid,
            'updatescript': self.updatescript,
            'port': PORT
            }
       
    
        print '> metadata: ', metadata
        with open(self.metadatafile, 'w') as f:
            json.dump(metadata, f, indent=4)
            
    def dump_geo(self):
        print '> dumping mesh'
        # save position
        with AtOrigin(self.node):
            cmds.file(self.geofile,
                        force=True,
                        pr=1,
                        typ="OBJexport",
                        exportSelected=True,
                        op="groups=0;materials=0;smoothing=0;normals=1")

        
    def edit_externally(self, path_to_app):
        blender_bridge_script = os.path.join(os.path.expanduser('~'), 'Documents', 'lab', 'm2b', 'b2m.py')
        blendfile = self.blendfile
        cmd = [path_to_app, '--python', blender_bridge_script, blendfile]
        print 'CMD', cmd
        subprocess.Popen(cmd)
        
# /// CONTROLS

def init():
    print 'init'
    selection = pm.selected()
    if len(selection) != 1:
        print 'Must select one object'
    else:
        selection = selection[0]
        uuid = cmds.ls(selection.name(), uuid=True)[0]
        bn = BridgedNode(uuid)
        bn.ingest()
        
        if bn.is_valid():
            if bn.has_history():
                print '> file has history. opening source.'
            else:
                bn.dump_geo()
                bn.dump_metadata()
        else:
            bn.dump_geo()
            bn.dump_metadata()


        open_port(6006)
        bn.edit_externally(get_blender_path())


def update(uuid):
    print '> running on', uuid
    bn = BridgedNode(uuid)
    bn.ingest()
    if bn.is_valid():
        with KeepSelection():
            print 'GEO', bn.geofile
            imported_node = load_geofile(bn.geofile)
            replace_mesh(bn.node, imported_node)
    else:
        print '> UUID invalid.'