From 66544273fd08fc22aedc4267b5716d0dbb8abd23 Mon Sep 17 00:00:00 2001
From: Ondrej Samohel <ondrej@samohel.cz>
Date: Fri, 1 Nov 2019 19:06:00 +0100
Subject: [PATCH] connect shapes to loaded Yeti Rig

---
 pype/plugins/maya/load/load_yeti_rig.py | 52 +++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 3 deletions(-)

diff --git a/pype/plugins/maya/load/load_yeti_rig.py b/pype/plugins/maya/load/load_yeti_rig.py
index eb75ff6bdc..a3e03e8a6c 100644
--- a/pype/plugins/maya/load/load_yeti_rig.py
+++ b/pype/plugins/maya/load/load_yeti_rig.py
@@ -1,9 +1,17 @@
-import pype.maya.plugin
 import os
+from collections import defaultdict
+
 from pypeapp import config
+import pype.maya.plugin
+from pype.maya import lib
 
 
 class YetiRigLoader(pype.maya.plugin.ReferenceLoader):
+    """
+    This loader will load Yeti rig. You can select something in scene and if it
+    has same ID as mesh published with rig, their shapes will be linked
+    together.
+    """
 
     families = ["yetiRig"]
     representations = ["ma"]
@@ -18,6 +26,32 @@ class YetiRigLoader(pype.maya.plugin.ReferenceLoader):
         import maya.cmds as cmds
         from avalon import maya
 
+        # get roots of selected hierarchies
+        selected_roots = []
+        for sel in cmds.ls(sl=True, long=True):
+            selected_roots.append(sel.split("|")[1])
+
+        # get all objects under those roots
+        selected_hierarchy = []
+        for root in selected_roots:
+            selected_hierarchy.append(cmds.listRelatives(
+                                      root,
+                                      allDescendents=True) or [])
+
+        # flatten the list and filter only shapes
+        shapes_flat = []
+        for root in selected_hierarchy:
+            shapes = cmds.ls(root, long=True, type="mesh") or []
+            for shape in shapes:
+                shapes_flat.append(shape)
+
+        # create dictionary of cbId and shape nodes
+        scene_lookup = defaultdict(list)
+        for node in shapes_flat:
+            cb_id = lib.get_id(node)
+            scene_lookup[cb_id] = node
+
+        # load rig
         with maya.maintained_selection():
             nodes = cmds.file(self.fname,
                               namespace=namespace,
@@ -26,6 +60,20 @@ class YetiRigLoader(pype.maya.plugin.ReferenceLoader):
                               groupReference=True,
                               groupName="{}:{}".format(namespace, name))
 
+        # for every shape node we've just loaded find matching shape by its
+        # cbId in selection. If found outMesh of scene shape will connect to
+        # inMesh of loaded shape.
+        for destination_node in nodes:
+            source_node = scene_lookup[lib.get_id(destination_node)]
+            if source_node:
+                self.log.info("found: {}".format(source_node))
+                self.log.info(
+                    "creating connection to {}".format(destination_node))
+
+                cmds.connectAttr("{}.outMesh".format(source_node),
+                                 "{}.inMesh".format(destination_node),
+                                 force=True)
+
         groupName = "{}:{}".format(namespace, name)
 
         presets = config.get_presets(project=os.environ['AVALON_PROJECT'])
@@ -38,6 +86,4 @@ class YetiRigLoader(pype.maya.plugin.ReferenceLoader):
                          c[0], c[1], c[2])
         self[:] = nodes
 
-        self.log.info("Yeti Rig Connection Manager will be available soon")
-
         return nodes