From 59d3ab14b2713c6853c85960f7731ac3f7ea829f Mon Sep 17 00:00:00 2001
From: Toke Stuart Jepsen <tokejepsen@gmail.com>
Date: Thu, 11 Jun 2020 23:38:55 +0100
Subject: [PATCH] Publish workfile

---
 .../harmony/publish/collect_workfile.py       | 28 +++++++++++++
 .../harmony/publish/extract_save_scene.py     | 13 ++++++
 .../harmony/publish/extract_workfile.py       | 34 ++++++++++++++++
 .../harmony/publish/increment_workfile.py     | 40 +++++++++++++++++++
 4 files changed, 115 insertions(+)
 create mode 100644 pype/plugins/harmony/publish/collect_workfile.py
 create mode 100644 pype/plugins/harmony/publish/extract_save_scene.py
 create mode 100644 pype/plugins/harmony/publish/extract_workfile.py
 create mode 100644 pype/plugins/harmony/publish/increment_workfile.py

diff --git a/pype/plugins/harmony/publish/collect_workfile.py b/pype/plugins/harmony/publish/collect_workfile.py
new file mode 100644
index 0000000000..7781eb0774
--- /dev/null
+++ b/pype/plugins/harmony/publish/collect_workfile.py
@@ -0,0 +1,28 @@
+import pyblish.api
+import os
+
+
+class CollectWorkfile(pyblish.api.ContextPlugin):
+    """Collect current script for publish."""
+
+    order = pyblish.api.CollectorOrder + 0.1
+    label = "Collect Workfile"
+    hosts = ["harmony"]
+
+    def process(self, context):
+        family = "workfile"
+        task = os.getenv("AVALON_TASK", None)
+        subset = family + task.capitalize()
+        basename = os.path.basename(context.data["currentFile"])
+
+        # Create instance
+        instance = context.create_instance(subset)
+        instance.data.update({
+            "subset": subset,
+            "label": basename,
+            "name": basename,
+            "family": family,
+            "families": [],
+            "representations": [],
+            "asset": os.environ["AVALON_ASSET"]
+        })
diff --git a/pype/plugins/harmony/publish/extract_save_scene.py b/pype/plugins/harmony/publish/extract_save_scene.py
new file mode 100644
index 0000000000..1733bdb95c
--- /dev/null
+++ b/pype/plugins/harmony/publish/extract_save_scene.py
@@ -0,0 +1,13 @@
+import pyblish.api
+from avalon import harmony
+
+
+class ExtractSaveScene(pyblish.api.ContextPlugin):
+    """Save scene for extraction."""
+
+    label = "Extract Save Scene"
+    order = pyblish.api.ExtractorOrder - 0.49
+    hosts = ["harmony"]
+
+    def process(self, instance):
+        harmony.save_scene()
diff --git a/pype/plugins/harmony/publish/extract_workfile.py b/pype/plugins/harmony/publish/extract_workfile.py
new file mode 100644
index 0000000000..3c51c94452
--- /dev/null
+++ b/pype/plugins/harmony/publish/extract_workfile.py
@@ -0,0 +1,34 @@
+import os
+import shutil
+
+import pype.api
+from avalon import harmony
+
+
+class ExtractWorkfile(pype.api.Extractor):
+    """Extract the connected nodes to the composite instance."""
+
+    label = "Extract Workfile"
+    hosts = ["harmony"]
+    families = ["workfile"]
+
+    def process(self, instance):
+        file_path = instance.context.data["currentFile"]
+        staging_dir = self.staging_dir(instance)
+
+        os.chdir(staging_dir)
+        shutil.make_archive(
+            instance.name,
+            "zip",
+            os.path.dirname(file_path)
+        )
+        zip_path = os.path.join(staging_dir, instance.name + ".zip")
+        self.log.info(f"Output zip file: {zip_path}")
+
+        representation = {
+            "name": "tpl",
+            "ext": "zip",
+            "files": "{}.zip".format(instance.name),
+            "stagingDir": staging_dir
+        }
+        instance.data["representations"] = [representation]
diff --git a/pype/plugins/harmony/publish/increment_workfile.py b/pype/plugins/harmony/publish/increment_workfile.py
new file mode 100644
index 0000000000..0b626ed190
--- /dev/null
+++ b/pype/plugins/harmony/publish/increment_workfile.py
@@ -0,0 +1,40 @@
+import os
+import shutil
+
+import pyblish.api
+from pype.action import get_errored_plugins_from_data
+from pype.lib import version_up
+import pype.api
+from avalon import harmony
+
+
+class IncrementWorkfile(pyblish.api.InstancePlugin):
+    """Increment the current workfile.
+
+    Saves the current scene with an increased version number.
+    """
+
+    label = "Increment Workfile"
+    order = pyblish.api.IntegratorOrder + 9.0
+    hosts = ["harmony"]
+    families = ["workfile"]
+    optional = True
+
+    def process(self, instance):
+        errored_plugins = get_errored_plugins_from_data(instance.context)
+        if errored_plugins:
+            raise RuntimeError(
+                "Skipping incrementing current file because submission to"
+                " deadline failed."
+            )
+
+        scene_dir = version_up(
+            os.path.dirname(instance.context.data["currentFile"])
+        )
+        scene_path = os.path.join(
+            scene_dir, os.path.basename(scene_dir) + ".xstage"
+        )
+
+        harmony.save_scene_as(scene_path)
+
+        self.log.info("Incremented workfile to: {}".format(scene_path))