From f419ed655ab0fbc76bbe8965d65c718dff2a7543 Mon Sep 17 00:00:00 2001
From: Jakub Jezek <jakub@orbi.tools>
Date: Fri, 9 Apr 2021 12:28:08 +0200
Subject: [PATCH] Hiero: adding new precollect otio plugins

---
 .../hiero/plugins/publish/extract_workfile.py |  50 +++
 .../plugins/publish/precollect_instances.py   | 316 ++++++++----------
 .../plugins/publish/precollect_workfile.py    |  76 ++---
 3 files changed, 210 insertions(+), 232 deletions(-)
 create mode 100644 openpype/hosts/hiero/plugins/publish/extract_workfile.py

diff --git a/openpype/hosts/hiero/plugins/publish/extract_workfile.py b/openpype/hosts/hiero/plugins/publish/extract_workfile.py
new file mode 100644
index 0000000000..e3d60465a2
--- /dev/null
+++ b/openpype/hosts/hiero/plugins/publish/extract_workfile.py
@@ -0,0 +1,50 @@
+import os
+import pyblish.api
+import openpype.api
+from openpype.hosts import resolve
+
+
+class ExtractWorkfile(openpype.api.Extractor):
+    """
+    Extractor export DRP workfile file representation
+    """
+
+    label = "Extract Workfile"
+    order = pyblish.api.ExtractorOrder
+    families = ["workfile"]
+    hosts = ["resolve"]
+
+    def process(self, instance):
+        # create representation data
+        if "representations" not in instance.data:
+            instance.data["representations"] = []
+
+        name = instance.data["name"]
+        project = instance.context.data["activeProject"]
+        staging_dir = self.staging_dir(instance)
+
+        resolve_workfile_ext = ".drp"
+        drp_file_name = name + resolve_workfile_ext
+        drp_file_path = os.path.normpath(
+            os.path.join(staging_dir, drp_file_name))
+
+        # write out the drp workfile
+        resolve.get_project_manager().ExportProject(
+            project.GetName(), drp_file_path)
+
+        # create drp workfile representation
+        representation_drp = {
+            'name': resolve_workfile_ext[1:],
+            'ext': resolve_workfile_ext[1:],
+            'files': drp_file_name,
+            "stagingDir": staging_dir,
+        }
+
+        instance.data["representations"].append(representation_drp)
+
+        # add sourcePath attribute to instance
+        if not instance.data.get("sourcePath"):
+            instance.data["sourcePath"] = drp_file_path
+
+        self.log.info("Added Resolve file representation: {}".format(
+            representation_drp))
diff --git a/openpype/hosts/hiero/plugins/publish/precollect_instances.py b/openpype/hosts/hiero/plugins/publish/precollect_instances.py
index 2b769afee1..242cfed254 100644
--- a/openpype/hosts/hiero/plugins/publish/precollect_instances.py
+++ b/openpype/hosts/hiero/plugins/publish/precollect_instances.py
@@ -1,224 +1,178 @@
-from compiler.ast import flatten
-from pyblish import api
+import pyblish
+import openpype
 from openpype.hosts.hiero import api as phiero
-import hiero
-# from openpype.hosts.hiero.api import lib
-# reload(lib)
-# reload(phiero)
+from openpype.hosts.hiero.api import lib
+from openpype.hosts.hiero.otio import hiero_export
+
+# # developer reload modules
+from pprint import pformat
+reload(lib)
+reload(phiero)
 
 
-class PreCollectInstances(api.ContextPlugin):
+class PrecollectInstances(pyblish.api.ContextPlugin):
     """Collect all Track items selection."""
 
-    order = api.CollectorOrder - 0.509
-    label = "Pre-collect Instances"
+    order = pyblish.api.CollectorOrder - 0.59
+    label = "Precollect Instances"
     hosts = ["hiero"]
 
-
     def process(self, context):
-        track_items = phiero.get_track_items(
-            selected=True, check_tagged=True, check_enabled=True)
-        # only return enabled track items
-        if not track_items:
-            track_items = phiero.get_track_items(
-                check_enabled=True, check_tagged=True)
-        # get sequence and video tracks
-        sequence = context.data["activeSequence"]
-        tracks = sequence.videoTracks()
-
-        # add collection to context
-        tracks_effect_items = self.collect_sub_track_items(tracks)
-
-        context.data["tracksEffectItems"] = tracks_effect_items
-
+        otio_timeline = context.data["otioTimeline"]
+        selected_timeline_items = phiero.get_track_items(
+                selected=True, check_enabled=True, check_tagged=True)
         self.log.info(
-            "Processing enabled track items: {}".format(len(track_items)))
+            "Processing enabled track items: {}".format(
+                selected_timeline_items))
 
-        for _ti in track_items:
-            data = {}
-            clip = _ti.source()
+        for track_item in selected_timeline_items:
 
-            # get clips subtracks and anotations
-            annotations = self.clip_annotations(clip)
-            subtracks = self.clip_subtrack(_ti)
-            self.log.debug("Annotations: {}".format(annotations))
-            self.log.debug(">> Subtracks: {}".format(subtracks))
+            data = dict()
 
-            # get pype tag data
-            tag_parsed_data = phiero.get_track_item_pype_data(_ti)
-            # self.log.debug(pformat(tag_parsed_data))
+            # get openpype tag data
+            tag_data = phiero.get_track_item_pype_data(track_item)
+            self.log.debug("__ tag_data: {}".format(pformat(tag_data)))
 
-            if not tag_parsed_data:
+            if not tag_data:
                 continue
 
-            if tag_parsed_data.get("id") != "pyblish.avalon.instance":
+            if tag_data.get("id") != "pyblish.avalon.instance":
                 continue
+
+            clip = track_item.source()
+
             # add tag data to instance data
             data.update({
-                k: v for k, v in tag_parsed_data.items()
+                k: v for k, v in tag_data.items()
                 if k not in ("id", "applieswhole", "label")
             })
 
-            asset = tag_parsed_data["asset"]
-            subset = tag_parsed_data["subset"]
-            review_track = tag_parsed_data.get("reviewTrack")
-            hiero_track = tag_parsed_data.get("heroTrack")
-            audio = tag_parsed_data.get("audio")
-
-            # remove audio attribute from data
-            data.pop("audio")
+            asset = tag_data["asset"]
+            subset = tag_data["subset"]
 
             # insert family into families
-            family = tag_parsed_data["family"]
-            families = [str(f) for f in tag_parsed_data["families"]]
+            family = tag_data["family"]
+            families = [str(f) for f in tag_data["families"]]
             families.insert(0, str(family))
 
-            track = _ti.parent()
-            media_source = _ti.source().mediaSource()
-            source_path = media_source.firstpath()
-            file_head = media_source.filenameHead()
-            file_info = media_source.fileinfos().pop()
-            source_first_frame = int(file_info.startFrame())
-
-            # apply only for review and master track instance
-            if review_track and hiero_track:
-                families += ["review", "ftrack"]
-
             data.update({
                 "name": "{} {} {}".format(asset, subset, families),
                 "asset": asset,
-                "item": _ti,
+                "item": track_item,
                 "families": families,
-
-                # tags
-                "tags": _ti.tags(),
-
-                # track item attributes
-                "track": track.name(),
-                "trackItem": track,
-                "reviewTrack": review_track,
-
-                # version data
-                "versionData": {
-                    "colorspace": _ti.sourceMediaColourTransform()
-                },
-
-                # source attribute
-                "source": source_path,
-                "sourceMedia": media_source,
-                "sourcePath": source_path,
-                "sourceFileHead": file_head,
-                "sourceFirst": source_first_frame,
-
-                # clip's effect
-                "clipEffectItems": subtracks
+                "publish": tag_data["publish"],
+                "fps": context.data["fps"]
             })
 
+            # otio clip data
+            otio_data = self.get_otio_clip_instance_data(
+                otio_timeline, track_item) or {}
+            self.log.debug("__ otio_data: {}".format(pformat(otio_data)))
+            data.update(otio_data)
+            self.log.debug("__ data: {}".format(pformat(data)))
+
+            # add resolution
+            self.get_resolution_to_data(data, context)
+
+            # create instance
             instance = context.create_instance(**data)
 
-            self.log.info("Creating instance.data: {}".format(instance.data))
+            # create shot instance for shot attributes create/update
+            self.create_shot_instance(context, **data)
 
-            if audio:
-                a_data = dict()
+            self.log.info("Creating instance: {}".format(instance))
+            self.log.debug(
+                "_ instance.data: {}".format(pformat(instance.data)))
 
-                # add tag data to instance data
-                a_data.update({
-                    k: v for k, v in tag_parsed_data.items()
-                    if k not in ("id", "applieswhole", "label")
-                })
+    def get_resolution_to_data(self, data, context):
+        assert data.get("otioClip"), "Missing `otioClip` data"
 
-                # create main attributes
-                subset = "audioMain"
-                family = "audio"
-                families = ["clip", "ftrack"]
-                families.insert(0, str(family))
+        # solve source resolution option
+        if data.get("sourceResolution", None):
+            otio_clip_metadata = data[
+                "otioClip"].media_reference.metadata
+            data.update({
+                "resolutionWidth": otio_clip_metadata[
+                        "openpype.source.width"],
+                "resolutionHeight": otio_clip_metadata[
+                        "openpype.source.height"],
+                "pixelAspect": otio_clip_metadata[
+                        "openpype.source.pixelAspect"]
+            })
+        else:
+            otio_tl_metadata = context.data["otioTimeline"].metadata
+            data.update({
+                "resolutionWidth": otio_tl_metadata["openpype.timeline.width"],
+                "resolutionHeight": otio_tl_metadata["openpype.timeline.height"],
+                "pixelAspect": otio_tl_metadata["openpype.timeline.pixelAspect"]
+            })
 
-                name = "{} {} {}".format(asset, subset, families)
+    def create_shot_instance(self, context, **data):
+        master_layer = data.get("heroTrack")
+        hierarchy_data = data.get("hierarchyData")
 
-                a_data.update({
-                    "name": name,
-                    "subset": subset,
-                    "asset": asset,
-                    "family": family,
-                    "families": families,
-                    "item": _ti,
+        if not master_layer:
+            return
 
-                    # tags
-                    "tags": _ti.tags(),
-                })
+        if not hierarchy_data:
+            return
 
-                a_instance = context.create_instance(**a_data)
-                self.log.info("Creating audio instance: {}".format(a_instance))
+        asset = data["asset"]
+        subset = "shotMain"
+
+        # insert family into families
+        family = "shot"
+
+        data.update({
+            "name": "{} {} {}".format(asset, subset, family),
+            "subset": subset,
+            "asset": asset,
+            "family": family,
+            "families": []
+        })
+
+        context.create_instance(**data)
+
+    def get_otio_clip_instance_data(self, otio_timeline, track_item):
+        """
+        Return otio objects for timeline, track and clip
+
+        Args:
+            timeline_item_data (dict): timeline_item_data from list returned by
+                                    resolve.get_current_timeline_items()
+            otio_timeline (otio.schema.Timeline): otio object
+
+        Returns:
+            dict: otio clip object
+
+        """
+        track_name = track_item.parent().name()
+        timeline_range = self.create_otio_time_range_from_timeline_item_data(
+            track_item)
+        for otio_clip in otio_timeline.each_clip():
+            track_name = otio_clip.parent().name
+            parent_range = otio_clip.range_in_parent()
+            if track_name not in track_name:
+                continue
+            if otio_clip.name not in track_item.name():
+                continue
+            if openpype.lib.is_overlapping_otio_ranges(
+                    parent_range, timeline_range, strict=True):
+
+                # add pypedata marker to otio_clip metadata
+                for marker in otio_clip.markers:
+                    if phiero.pype_tag_name in marker.name:
+                        otio_clip.metadata.update(marker.metadata)
+                return {"otioClip": otio_clip}
+
+        return None
 
     @staticmethod
-    def clip_annotations(clip):
-        """
-        Returns list of Clip's hiero.core.Annotation
-        """
-        annotations = []
-        subTrackItems = flatten(clip.subTrackItems())
-        annotations += [item for item in subTrackItems if isinstance(
-            item, hiero.core.Annotation)]
-        return annotations
+    def create_otio_time_range_from_timeline_item_data(track_item):
+        timeline = phiero.get_current_sequence()
+        frame_start = int(track_item.timelineIn())
+        frame_duration = int(track_item.sourceDuration())
+        fps = timeline.framerate().toFloat()
 
-    @staticmethod
-    def clip_subtrack(clip):
-        """
-        Returns list of Clip's hiero.core.SubTrackItem
-        """
-        subtracks = []
-        subTrackItems = flatten(clip.parent().subTrackItems())
-        for item in subTrackItems:
-            # avoid all anotation
-            if isinstance(item, hiero.core.Annotation):
-                continue
-            # # avoid all not anaibled
-            if not item.isEnabled():
-                continue
-            subtracks.append(item)
-        return subtracks
-
-    @staticmethod
-    def collect_sub_track_items(tracks):
-        """
-        Returns dictionary with track index as key and list of subtracks
-        """
-        # collect all subtrack items
-        sub_track_items = dict()
-        for track in tracks:
-            items = track.items()
-
-            # skip if no clips on track > need track with effect only
-            if items:
-                continue
-
-            # skip all disabled tracks
-            if not track.isEnabled():
-                continue
-
-            track_index = track.trackIndex()
-            _sub_track_items = flatten(track.subTrackItems())
-
-            # continue only if any subtrack items are collected
-            if len(_sub_track_items) < 1:
-                continue
-
-            enabled_sti = list()
-            # loop all found subtrack items and check if they are enabled
-            for _sti in _sub_track_items:
-                # checking if not enabled
-                if not _sti.isEnabled():
-                    continue
-                if isinstance(_sti, hiero.core.Annotation):
-                    continue
-                # collect the subtrack item
-                enabled_sti.append(_sti)
-
-            # continue only if any subtrack items are collected
-            if len(enabled_sti) < 1:
-                continue
-
-            # add collection of subtrackitems to dict
-            sub_track_items[track_index] = enabled_sti
-
-        return sub_track_items
+        return hiero_export.create_otio_time_range(
+            frame_start, frame_duration, fps)
diff --git a/openpype/hosts/hiero/plugins/publish/precollect_workfile.py b/openpype/hosts/hiero/plugins/publish/precollect_workfile.py
index ef7d07421b..22201cafe3 100644
--- a/openpype/hosts/hiero/plugins/publish/precollect_workfile.py
+++ b/openpype/hosts/hiero/plugins/publish/precollect_workfile.py
@@ -1,74 +1,48 @@
-import os
 import pyblish.api
+import hiero.ui
 from openpype.hosts.hiero import api as phiero
 from avalon import api as avalon
+from pprint import pformat
+from openpype.hosts.hiero.otio import hiero_export
 
 
-class PreCollectWorkfile(pyblish.api.ContextPlugin):
+class PrecollectWorkfile(pyblish.api.ContextPlugin):
     """Inject the current working file into context"""
 
-    label = "Pre-collect Workfile"
-    order = pyblish.api.CollectorOrder - 0.51
+    label = "Precollect Workfile"
+    order = pyblish.api.CollectorOrder - 0.6
 
     def process(self, context):
+
         asset = avalon.Session["AVALON_ASSET"]
         subset = "workfile"
-
         project = phiero.get_current_project()
-        active_sequence = phiero.get_current_sequence()
-        video_tracks = active_sequence.videoTracks()
-        audio_tracks = active_sequence.audioTracks()
-        current_file = project.path()
-        staging_dir = os.path.dirname(current_file)
-        base_name = os.path.basename(current_file)
+        active_timeline = hiero.ui.activeSequence()
+        fps = active_timeline.framerate().toFloat()
 
-        # get workfile's colorspace properties
-        _clrs = {}
-        _clrs["useOCIOEnvironmentOverride"] =  project.useOCIOEnvironmentOverride()  # noqa
-        _clrs["lutSetting16Bit"] = project.lutSetting16Bit()
-        _clrs["lutSetting8Bit"] = project.lutSetting8Bit()
-        _clrs["lutSettingFloat"] = project.lutSettingFloat()
-        _clrs["lutSettingLog"] = project.lutSettingLog()
-        _clrs["lutSettingViewer"] = project.lutSettingViewer()
-        _clrs["lutSettingWorkingSpace"] = project.lutSettingWorkingSpace()
-        _clrs["lutUseOCIOForExport"] = project.lutUseOCIOForExport()
-        _clrs["ocioConfigName"] = project.ocioConfigName()
-        _clrs["ocioConfigPath"] = project.ocioConfigPath()
-
-        # set main project attributes to context
-        context.data["activeProject"] = project
-        context.data["activeSequence"] = active_sequence
-        context.data["videoTracks"] = video_tracks
-        context.data["audioTracks"] = audio_tracks
-        context.data["currentFile"] = current_file
-        context.data["colorspace"] = _clrs
-
-        self.log.info("currentFile: {}".format(current_file))
-
-        # creating workfile representation
-        representation = {
-            'name': 'hrox',
-            'ext': 'hrox',
-            'files': base_name,
-            "stagingDir": staging_dir,
-        }
+        # adding otio timeline to context
+        otio_timeline = hiero_export.create_otio_timeline()
 
         instance_data = {
             "name": "{}_{}".format(asset, subset),
             "asset": asset,
             "subset": "{}{}".format(asset, subset.capitalize()),
             "item": project,
-            "family": "workfile",
-
-            # version data
-            "versionData": {
-                "colorspace": _clrs
-            },
-
-            # source attribute
-            "sourcePath": current_file,
-            "representations": [representation]
+            "family": "workfile"
         }
 
+        # create instance with workfile
         instance = context.create_instance(**instance_data)
+
+        # update context with main project attributes
+        context_data = {
+            "activeProject": project,
+            "otioTimeline": otio_timeline,
+            "currentFile": project.path(),
+            "fps": fps,
+        }
+        context.data.update(context_data)
+
         self.log.info("Creating instance: {}".format(instance))
+        self.log.debug("__ instance.data: {}".format(pformat(instance.data)))
+        self.log.debug("__ context_data: {}".format(pformat(context_data)))