From f43253c773f1c8a64d74564ba84c3000c3e64fee Mon Sep 17 00:00:00 2001
From: wijnand <wikoreman@gmail.com>
Date: Thu, 12 Jul 2018 14:26:10 +0200
Subject: [PATCH] added collect and validate aovs for extend logic

---
 .../maya/publish/collect_render_layer_aovs.py | 91 +++++++++++++++++++
 .../maya/publish/validate_renderlayer_aovs.py | 60 ++++++++++++
 2 files changed, 151 insertions(+)
 create mode 100644 colorbleed/plugins/maya/publish/collect_render_layer_aovs.py
 create mode 100644 colorbleed/plugins/maya/publish/validate_renderlayer_aovs.py

diff --git a/colorbleed/plugins/maya/publish/collect_render_layer_aovs.py b/colorbleed/plugins/maya/publish/collect_render_layer_aovs.py
new file mode 100644
index 0000000000..64adb11258
--- /dev/null
+++ b/colorbleed/plugins/maya/publish/collect_render_layer_aovs.py
@@ -0,0 +1,91 @@
+from maya import cmds
+
+import pyblish.api
+
+import colorbleed.maya.lib as lib
+
+
+class CollectRenderLayerAOVS(pyblish.api.InstancePlugin):
+    """Validate all render layer's AOVs / Render Elements are registered in
+    the database
+
+    This validator is important to be able to Extend Frames
+
+    Technical information:
+    Each renderer uses different logic to work with render passes.
+    VRay - RenderElement
+        Simple node connection to the actual renderLayer node
+
+    Arnold - AOV:
+        Uses its own render settings node and connects an aiOAV to it
+
+    Redshift - AOV:
+        Uses its own render settings node and RedshiftAOV node. It is not
+        connected but all AOVs are enabled for all render layers by default.
+
+    """
+
+    order = pyblish.api.CollectorOrder + 0.01
+    label = "Render Elements / AOVs"
+    hosts = ["maya"]
+    families = ["colorbleed.renderlayer"]
+
+    def process(self, instance):
+
+        # Get renderer
+        renderer = cmds.getAttr("defaultRenderGlobals.currentRenderer")
+
+        self.log.info("Renderer found: {}".format(renderer))
+
+        rp_node_types = {"vray": "VRayRenderElement",
+                         "arnold": "aiAOV",
+                         "redshift": "RedshiftAOV"}
+
+        if renderer not in rp_node_types.keys():
+            self.log.error("Unsupported renderer found: '{}'".format(renderer))
+            return
+
+        result = []
+
+        # Collect all AOVs / Render Elements
+        with lib.renderlayer(instance.name):
+
+            node_type = rp_node_types[renderer]
+            render_elements = cmds.ls(type=node_type)
+
+            # Check if AOVs / Render Elements are enabled
+            for element in render_elements:
+                enabled = cmds.getAttr("{}.enabled".format(element))
+                if not enabled:
+                    continue
+
+                pass_name = self.get_pass_name(renderer, element)
+                render_pass = "%s.%s" % (instance.name, pass_name)
+
+                result.append(render_pass)
+
+        self.log.info("Found {} render elements / AOVs for "
+                      "'{}'".format(len(result), instance.name))
+
+        instance.data["renderPasses"] = result
+
+    def get_pass_name(self, renderer, node):
+
+        if renderer == "vray":
+            vray_node_attr = next(attr for attr in cmds.listAttr(node)
+                                  if attr.startswith("vray_name"))
+
+            pass_type = vray_node_attr.rsplit("_", 1)[-1]
+            if pass_type == "extratex":
+                vray_node_attr = "vray_explicit_name_extratex"
+
+            # Node type is in the attribute name but we need to check if value
+            # of the attribute as it can be changed
+            pass_name = cmds.getAttr("{}.{}".format(node, vray_node_attr))
+
+        elif renderer in ["arnold", "redshift"]:
+            pass_name = cmds.getAttr("{}.name".format(node))
+        else:
+            raise RuntimeError("Unsupported renderer: '{}'".format(renderer))
+
+        return pass_name
\ No newline at end of file
diff --git a/colorbleed/plugins/maya/publish/validate_renderlayer_aovs.py b/colorbleed/plugins/maya/publish/validate_renderlayer_aovs.py
new file mode 100644
index 0000000000..de54be2147
--- /dev/null
+++ b/colorbleed/plugins/maya/publish/validate_renderlayer_aovs.py
@@ -0,0 +1,60 @@
+import pyblish.api
+
+from avalon import io, api
+import colorbleed.api
+
+
+class ValidateRenderLayerAOVs(pyblish.api.InstancePlugin):
+    """Validate created AOVs / RenderElement is registered in the database
+
+    Each render element is registered as a subset which is formatted based on
+    the render layer and the render element, example:
+
+        <render layer>.<render element>
+
+    This translates to something like this:
+
+        CHAR.diffuse
+
+    This check is only used when the Extend Frames in 'renderGlobalsDefault'
+    is toggled on.
+
+    This check is needed to ensure the render output is still complete
+
+    """
+
+    order = pyblish.api.ValidatorOrder + 0.1
+    label = "Render Passes / AOVs Are Registered"
+    hosts = ["maya"]
+    families = ["colorbleed.renderlayer"]
+    actions = [colorbleed.api.SelectInvalidAction]
+
+    def process(self, instance):
+        invalid = self.get_invalid(instance)
+        if invalid:
+            raise RuntimeError("Found unregistered subsets: {}".format(invalid))
+
+    def get_invalid(self, instance):
+
+        invalid = []
+
+        asset_name = api.Session["AVALON_ASSET"]
+        render_passses = instance.data.get("renderPasses", [])
+        for render_pass in render_passses:
+
+            is_valid = self.validate_subset_registered(asset_name, render_pass)
+            if not is_valid:
+                invalid.append(render_pass)
+
+        return invalid
+
+    def validate_subset_registered(self, asset_name, subset_name):
+        """Check if subset is registered in the database under the asset"""
+
+        asset = io.find_one({"type": "asset", "name": asset_name})
+        is_valid = io.find_one({"type": "subset",
+                                "name": subset_name,
+                                "parent": asset["_id"]})
+
+        return is_valid
+