From a39626d71040afb30182a3c5c0e45e5c931bb49c Mon Sep 17 00:00:00 2001
From: Jakub Jezek <jakubjezek001@gmail.com>
Date: Fri, 4 Aug 2023 14:57:20 +0200
Subject: [PATCH] refactor collect_writes

- also update exctract slates so it expect updated data
- and move some data to collect instance data
- also renaming collect instance data to host related name
---
 ..._data.py => collect_nuke_instance_data.py} |  17 +-
 .../nuke/plugins/publish/collect_writes.py    | 412 ++++++++++++------
 .../plugins/publish/extract_slate_frame.py    |   2 +-
 3 files changed, 297 insertions(+), 134 deletions(-)
 rename openpype/hosts/nuke/plugins/publish/{collect_instance_data.py => collect_nuke_instance_data.py} (71%)

diff --git a/openpype/hosts/nuke/plugins/publish/collect_instance_data.py b/openpype/hosts/nuke/plugins/publish/collect_nuke_instance_data.py
similarity index 71%
rename from openpype/hosts/nuke/plugins/publish/collect_instance_data.py
rename to openpype/hosts/nuke/plugins/publish/collect_nuke_instance_data.py
index 3908aef4bc..ffcf49d697 100644
--- a/openpype/hosts/nuke/plugins/publish/collect_instance_data.py
+++ b/openpype/hosts/nuke/plugins/publish/collect_nuke_instance_data.py
@@ -2,11 +2,13 @@ import nuke
 import pyblish.api
 
 
-class CollectInstanceData(pyblish.api.InstancePlugin):
-    """Collect all nodes with Avalon knob."""
+class CollectNukeInstanceData(pyblish.api.InstancePlugin):
+    """Collect Nuke instance data
+
+    """
 
     order = pyblish.api.CollectorOrder - 0.49
-    label = "Collect Instance Data"
+    label = "Collect Nuke Instance Data"
     hosts = ["nuke", "nukeassist"]
 
     # presets
@@ -40,5 +42,14 @@ class CollectInstanceData(pyblish.api.InstancePlugin):
             "pixelAspect": pixel_aspect
 
         })
+
+        # add review family if review activated on instance
+        if instance.data.get("review"):
+            instance.data["families"].append("review")
+
+        # add creator attributes to instance
+        creator_attributes = instance.data["creator_attributes"]
+        instance.data.update(creator_attributes)
+
         self.log.debug("Collected instance: {}".format(
             instance.data))
diff --git a/openpype/hosts/nuke/plugins/publish/collect_writes.py b/openpype/hosts/nuke/plugins/publish/collect_writes.py
index 0d552b7381..307d323908 100644
--- a/openpype/hosts/nuke/plugins/publish/collect_writes.py
+++ b/openpype/hosts/nuke/plugins/publish/collect_writes.py
@@ -1,5 +1,4 @@
 import os
-from pprint import pformat
 import nuke
 import pyblish.api
 from openpype.hosts.nuke import api as napi
@@ -15,33 +14,16 @@ class CollectNukeWrites(pyblish.api.InstancePlugin,
     hosts = ["nuke", "nukeassist"]
     families = ["render", "prerender", "image"]
 
+    # cashing
+    _write_nodes = {}
+    _frame_ranges = {}
+
     def process(self, instance):
-        self.log.debug(pformat(instance.data))
-        creator_attributes = instance.data["creator_attributes"]
-        instance.data.update(creator_attributes)
 
         group_node = instance.data["transientData"]["node"]
         render_target = instance.data["render_target"]
-        family = instance.data["family"]
-        families = instance.data["families"]
 
-        # add targeted family to families
-        instance.data["families"].append(
-            "{}.{}".format(family, render_target)
-        )
-        self.log.debug("Appending render target to families: {}.{}".format(
-            family, render_target)
-        )
-        if instance.data.get("review"):
-            instance.data["families"].append("review")
-
-        child_nodes = napi.get_instance_group_node_childs(instance)
-        instance.data["transientData"]["childNodes"] = child_nodes
-
-        write_node = None
-        for x in child_nodes:
-            if x.Class() == "Write":
-                write_node = x
+        write_node = self._write_node_helper(instance)
 
         if write_node is None:
             self.log.warning(
@@ -51,131 +33,134 @@ class CollectNukeWrites(pyblish.api.InstancePlugin,
             )
             return
 
-        instance.data["writeNode"] = write_node
-        self.log.debug("checking instance: {}".format(instance))
+        # get colorspace and add to version data
+        colorspace = napi.get_colorspace_from_node(write_node)
 
-        # Determine defined file type
-        ext = write_node["file_type"].value()
+        if render_target == "frames":
+            self._set_existing_files_data(instance, colorspace)
 
-        # Get frame range
-        handle_start = instance.context.data["handleStart"]
-        handle_end = instance.context.data["handleEnd"]
-        first_frame = int(nuke.root()["first_frame"].getValue())
-        last_frame = int(nuke.root()["last_frame"].getValue())
-        frame_length = int(last_frame - first_frame + 1)
+        elif render_target == "frames_farm":
+            collected_frames = self._set_existing_files_data(
+                instance, colorspace)
 
-        if write_node["use_limit"].getValue():
-            first_frame = int(write_node["first"].getValue())
-            last_frame = int(write_node["last"].getValue())
+            self._set_expected_files(instance, collected_frames)
+
+            self._add_farm_instance_data(instance)
+
+        elif render_target == "farm":
+            self._add_farm_instance_data(instance)
+
+        # set additional instance data
+        self._set_additional_instance_data(instance, render_target, colorspace)
+
+    def _set_existing_files_data(self, instance, colorspace):
+        """Set existing files data to instance data.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+            colorspace (str): colorspace
+
+        Returns:
+            list: collected frames
+        """
+        collected_frames = self._get_collected_frames(instance)
+
+        representation = self._get_existing_frames_representation(
+            instance, collected_frames
+        )
+
+        # inject colorspace data
+        self.set_representation_colorspace(
+            representation, instance.context,
+            colorspace=colorspace
+        )
+
+        instance.data["representations"].append(representation)
+
+        return collected_frames
+
+    def _set_expected_files(self, instance, collected_frames):
+        """Set expected files to instance data.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+            collected_frames (list): collected frames
+        """
+        write_node = self._write_node_helper(instance)
 
         write_file_path = nuke.filename(write_node)
         output_dir = os.path.dirname(write_file_path)
 
-        # get colorspace and add to version data
-        colorspace = napi.get_colorspace_from_node(write_node)
+        instance.data["expectedFiles"] = [
+            os.path.join(output_dir, source_file)
+            for source_file in collected_frames
+        ]
 
-        self.log.debug('output dir: {}'.format(output_dir))
+    def _get_frame_range_data(self, instance):
+        """Get frame range data from instance.
 
-        if render_target in ["frames", "frames_farm"]:
-            representation = {
-                'name': ext,
-                'ext': ext,
-                "stagingDir": output_dir,
-                "tags": []
-            }
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
 
-            # get file path knob
-            node_file_knob = write_node["file"]
-            # list file paths based on input frames
-            expected_paths = list(sorted({
-                node_file_knob.evaluate(frame)
-                for frame in range(first_frame, last_frame + 1)
-            }))
+        Returns:
+            tuple: first_frame, last_frame
+        """
 
-            # convert only to base names
-            expected_filenames = [
-                os.path.basename(filepath)
-                for filepath in expected_paths
-            ]
+        instance_name = instance.data["name"]
 
-            # make sure files are existing at folder
-            collected_frames = [
-                filename
-                for filename in os.listdir(output_dir)
-                if filename in expected_filenames
-            ]
+        if self._frame_ranges.get(instance_name):
+            # return cashed write node
+            return self._frame_ranges[instance_name]
 
-            if collected_frames:
-                collected_frames_len = len(collected_frames)
-                frame_start_str = "%0{}d".format(
-                    len(str(last_frame))) % first_frame
-                representation['frameStart'] = frame_start_str
+        write_node = self._write_node_helper(instance)
 
-                # in case slate is expected and not yet rendered
-                self.log.debug("_ frame_length: {}".format(frame_length))
-                self.log.debug("_ collected_frames_len: {}".format(
-                    collected_frames_len))
+        # Get frame range from workfile
+        first_frame = int(nuke.root()["first_frame"].getValue())
+        last_frame = int(nuke.root()["last_frame"].getValue())
 
-                # this will only run if slate frame is not already
-                # rendered from previews publishes
-                if (
-                    "slate" in families
-                    and frame_length == collected_frames_len
-                    and family == "render"
-                ):
-                    frame_slate_str = (
-                        "{{:0{}d}}".format(len(str(last_frame)))
-                    ).format(first_frame - 1)
+        # Get frame range from write node if activated
+        if write_node["use_limit"].getValue():
+            first_frame = int(write_node["first"].getValue())
+            last_frame = int(write_node["last"].getValue())
 
-                    slate_frame = collected_frames[0].replace(
-                        frame_start_str, frame_slate_str)
-                    collected_frames.insert(0, slate_frame)
+        # add to cache
+        self._frame_ranges[instance_name] = (first_frame, last_frame)
 
-                if collected_frames_len == 1:
-                    representation['files'] = collected_frames.pop()
-                else:
-                    representation['files'] = collected_frames
+        return first_frame, last_frame
 
-                # inject colorspace data
-                self.set_representation_colorspace(
-                    representation, instance.context,
-                    colorspace=colorspace
-                )
+    def _set_additional_instance_data(
+        self, instance, render_target, colorspace
+    ):
+        """Set additional instance data.
 
-            instance.data["representations"].append(representation)
-            self.log.info("Publishing rendered frames ...")
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+            render_target (str): render target
+            colorspace (str): colorspace
+        """
+        family = instance.data["family"]
 
-            if render_target == "frames_farm":
-                # Farm rendering
-                instance.data["toBeRenderedOn"] = "deadline"
-                instance.data["transfer"] = False
-                instance.data["farm"] = True  # to skip integrate
-                self.log.info("Farm rendering ON ...")
+        # add targeted family to families
+        instance.data["families"].append(
+            "{}.{}".format(family, render_target)
+        )
+        self.log.info("Appending render target to families: {}.{}".format(
+            family, render_target)
+        )
 
-                self.log.info(
-                    "Adding collected files %s to expectedFiles instance.data",
-                    collected_frames
-                )
-                if "expectedFiles" not in instance.data:
-                    instance.data["expectedFiles"] = list()
-                    for source_file in collected_frames:
-                        instance.data["expectedFiles"].append(
-                            os.path.join(output_dir, source_file)
-                        )
+        write_node = self._write_node_helper(instance)
 
-        elif render_target == "farm":
-            farm_keys = ["farm_chunk", "farm_priority", "farm_concurrency"]
-            for key in farm_keys:
-                # Skip if key is not in creator attributes
-                if key not in creator_attributes:
-                    continue
-                # Add farm attributes to instance
-                instance.data[key] = creator_attributes[key]
+        # Determine defined file type
+        ext = write_node["file_type"].value()
 
-            # Farm rendering
-            instance.data["transfer"] = False
-            instance.data["farm"] = True
-            self.log.info("Farm rendering ON ...")
+        # get frame range data
+        handle_start = instance.context.data["handleStart"]
+        handle_end = instance.context.data["handleEnd"]
+        first_frame, last_frame = self._get_frame_range_data(instance)
+
+        # get output paths
+        write_file_path = nuke.filename(write_node)
+        output_dir = os.path.dirname(write_file_path)
 
         # TODO: remove this when we have proper colorspace support
         version_data = {
@@ -209,10 +194,6 @@ class CollectNukeWrites(pyblish.api.InstancePlugin,
                 "frameEndHandle": last_frame,
             })
 
-        # make sure rendered sequence on farm will
-        # be used for extract review
-        if not instance.data.get("review"):
-            instance.data["useSequenceForReview"] = False
 
         # TODO temporarily set stagingDir as persistent for backward
         # compatibility. This is mainly focused on `renders`folders which
@@ -220,4 +201,175 @@ class CollectNukeWrites(pyblish.api.InstancePlugin,
         # this logic should be removed and replaced with custom staging dir
         instance.data["stagingDir_persistent"] = True
 
-        self.log.debug("instance.data: {}".format(pformat(instance.data)))
+    def _write_node_helper(self, instance):
+        """Helper function to get write node from instance.
+
+        Also sets instance transient data with child nodes.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+
+        Returns:
+            nuke.Node: write node
+        """
+        instance_name = instance.data["name"]
+
+        if self._write_nodes.get(instance_name):
+            # return cashed write node
+            return self._write_nodes[instance_name]
+
+        # get all child nodes from group node
+        child_nodes = napi.get_instance_group_node_childs(instance)
+
+        # set child nodes to instance transient data
+        instance.data["transientData"]["childNodes"] = child_nodes
+
+        write_node = None
+        for node_ in child_nodes:
+            if node_.Class() == "Write":
+                write_node = node_
+
+        if write_node:
+            # for slate frame extraction
+            instance.data["transientData"]["writeNode"] = write_node
+            # add to cache
+            self._write_nodes[instance_name] = write_node
+
+            return self._write_nodes[instance_name]
+
+    def _get_existing_frames_representation(
+        self,
+        instance,
+        collected_frames
+    ):
+        """Get existing frames representation.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+            collected_frames (list): collected frames
+
+        Returns:
+            dict: representation
+        """
+
+        first_frame, last_frame = self._get_frame_range_data(instance)
+
+        write_node = self._write_node_helper(instance)
+
+        write_file_path = nuke.filename(write_node)
+        output_dir = os.path.dirname(write_file_path)
+
+        # Determine defined file type
+        ext = write_node["file_type"].value()
+
+        representation = {
+            "name": ext,
+            "ext": ext,
+            "stagingDir": output_dir,
+            "tags": []
+        }
+
+        frame_start_str = "%0{}d".format(
+            len(str(last_frame))) % first_frame
+        representation['frameStart'] = frame_start_str
+
+        # set slate frame
+        collected_frames = self._add_slate_frame_to_collected_frames(
+            instance,
+            collected_frames,
+            frame_start_str,
+            first_frame,
+            last_frame
+        )
+
+        if len(collected_frames) == 1:
+            representation['files'] = collected_frames.pop()
+        else:
+            representation['files'] = collected_frames
+
+        return representation
+
+    def _add_slate_frame_to_collected_frames(
+        self,
+        instance,
+        collected_frames,
+        frame_start_str,
+        first_frame,
+        last_frame
+    ):
+        """Set slate frame."""
+        frame_length = int(last_frame - first_frame + 1)
+
+        # this will only run if slate frame is not already
+        # rendered from previews publishes
+        if (
+            "slate" in instance.data["families"]
+            and frame_length == len(collected_frames)
+            and instance.data["family"] == "render"
+        ):
+            frame_slate_str = (
+                "{{:0{}d}}".format(len(str(last_frame)))
+            ).format(first_frame - 1)
+
+        slate_frame = collected_frames[0].replace(
+            frame_start_str, frame_slate_str)
+        collected_frames.insert(0, slate_frame)
+
+        return collected_frames
+
+    def _add_farm_instance_data(self, instance):
+        """Add farm publishing related instance data.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+        """
+
+        # make sure rendered sequence on farm will
+        # be used for extract review
+        if not instance.data.get("review"):
+            instance.data["useSequenceForReview"] = False
+
+        # Farm rendering
+        instance.data["transfer"] = False
+        instance.data["farm"] = True
+        self.log.info("Farm rendering ON ...")
+
+    def _get_collected_frames(self, instance):
+        """Get collected frames.
+
+        Args:
+            instance (pyblish.api.Instance): pyblish instance
+
+        Returns:
+            list: collected frames
+        """
+
+        first_frame, last_frame = self._get_frame_range_data(instance)
+
+        write_node = self._write_node_helper(instance)
+
+        write_file_path = nuke.filename(write_node)
+        output_dir = os.path.dirname(write_file_path)
+
+        # get file path knob
+        node_file_knob = write_node["file"]
+        # list file paths based on input frames
+        expected_paths = list(sorted({
+            node_file_knob.evaluate(frame)
+            for frame in range(first_frame, last_frame + 1)
+        }))
+
+        # convert only to base names
+        expected_filenames = [
+            os.path.basename(filepath)
+            for filepath in expected_paths
+        ]
+
+        # make sure files are existing at folder
+        collected_frames = [
+            filename
+            for filename in os.listdir(output_dir)
+            if filename in expected_filenames
+        ]
+
+        return collected_frames
diff --git a/openpype/hosts/nuke/plugins/publish/extract_slate_frame.py b/openpype/hosts/nuke/plugins/publish/extract_slate_frame.py
index 06c086b10d..25262a7418 100644
--- a/openpype/hosts/nuke/plugins/publish/extract_slate_frame.py
+++ b/openpype/hosts/nuke/plugins/publish/extract_slate_frame.py
@@ -249,7 +249,7 @@ class ExtractSlateFrame(publish.Extractor):
 
         # Add file to representation files
         # - get write node
-        write_node = instance.data["writeNode"]
+        write_node = instance.data["transientData"]["writeNode"]
         # - evaluate filepaths for first frame and slate frame
         first_filename = os.path.basename(
             write_node["file"].evaluate(first_frame))