Add collector custom arguments

hooks/collector.py

@@ -149,7 +149,7 @@ def settings(self):
         """
         return {}
 
-    def process_current_session(self, settings, parent_item):
+    def process_current_session(self, settings, parent_item, collection_args=None):
         """
         Analyzes the current scene open in a DCC and parents a subtree of items
         under the parent_item passed in.
@@ -161,7 +161,7 @@ def process_current_session(self, settings, parent_item):
         # default implementation does not do anything
         pass
 
-    def process_file(self, settings, parent_item, path):
+    def process_file(self, settings, parent_item, path, collection_args=None):
         """
         Analyzes the given file and creates one or more items
         to represent it.

python/tk_multi_publish2/api/manager.py

@@ -93,7 +93,7 @@ def __init__(self, publish_logger=None):
             base_class=self._bundle.base_hooks.PostPhaseHook
         )
 
-    def collect_files(self, file_paths):
+    def collect_files(self, file_paths, collection_args=None):
         """
         Run the collection logic to populate the publish tree with items for
         each supplied path.
@@ -127,7 +127,8 @@ def collect_files(self, file_paths):
                 # that are collected.
                 self._collector_instance.run_process_file(
                     self.tree.root_item,
-                    file_path
+                    file_path,
+                    collection_args,
                 )
 
             # get a list of all items in the tree after collection
@@ -156,7 +157,7 @@ def collect_files(self, file_paths):
 
         return new_items
 
-    def collect_session(self):
+    def collect_session(self, collection_args=None):
         """
         Run the collection logic to populate the tree with items to publish.
 
@@ -180,7 +181,7 @@ def collect_session(self):
         # we supply the root item of the tree for parenting of items that
         # are collected.
         self._collector_instance.run_process_current_session(
-            self.tree.root_item)
+            self.tree.root_item, collection_args)
 
         # get a list of all items in the tree after collection
         items_after = list(self.tree)

python/tk_multi_publish2/api/plugins/collector_instance.py

@@ -8,6 +8,7 @@
 # agreement to the Shotgun Pipeline Toolkit Source Code License. All rights
 # not expressly granted therein are reserved by Shotgun Software Inc.
 
+import inspect
 import traceback
 
 import sgtk
@@ -38,50 +39,73 @@ def _create_hook_instance(self, path):
         plugin.id = path
         return plugin
 
-    def run_process_file(self, item, path):
+    def _get_process_args(self, process_method, needed_args, new_supported_args):
+        """
+        Find args to be passed for the given process method
+        """
+        args = []
+        if hasattr(self._hook_instance.__class__, "settings"):
+            # this hook has a 'settings' property defined. it is expecting
+            # 'settings' to be passed to the processing method.
+            args.append(self.settings)
+
+        # needed_args are added in any case
+        args.extend(needed_args)
+
+        # for backward compatibility
+        # now check if the collector hook supports new arguments added afterward
+        hook_process_args = inspect.getargspec(process_method).args
+        for arg_name, arg_value in new_supported_args:
+            has_optional_arg_set = arg_name in hook_process_args
+            if has_optional_arg_set:
+                # this hook has the given argument.
+                # it is expecting to be passed to the processing method.
+                args.append(arg_value)
+
+        return args
+
+    def run_process_file(self, item, path, collection_args=None):
         """
         Executes the hook process_file method
 
         :param item: Item to parent collected items under.
         :param path: The path of the file to collect
+        :param collection_args: Custom data to be passed to the collector hook
 
         :returns: None (item creation handles parenting)
         """
+        # build args to be passed to process method
+        needed_args = [item, path]
+        new_supported_args = [('collection_args', collection_args)]
+        process_args = self._get_process_args(self._hook_instance.process_file, needed_args, new_supported_args)
+
         try:
-            if hasattr(self._hook_instance.__class__, "settings"):
-                # this hook has a 'settings' property defined. it is expecting
-                # 'settings' to be passed to the processing method.
-                return self._hook_instance.process_file(
-                    self.settings, item, path)
-            else:
-                # the hook hasn't been updated to handle collector settings.
-                # call the method without a settings argument
-                return self._hook_instance.process_file(item, path)
+            return self._hook_instance.process_file(*process_args)
         except Exception:
             error_msg = traceback.format_exc()
             logger.error(
                 "Error running process_file for %s. %s" %
                 (self, error_msg)
             )
 
-    def run_process_current_session(self, item):
+    def run_process_current_session(self, item, collection_args=None):
         """
         Executes the hook process_current_session method
 
         :param item: Item to parent collected items under.
+        :param collection_args: Custom data to be passed to the collector hook
 
         :returns: None (item creation handles parenting)
         """
+        # build args to be passed to process method
+        needed_args = [item]
+        new_supported_args = [('collection_args', collection_args)]
+        process_args = self._get_process_args(self._hook_instance.process_current_session,
+                                              needed_args,
+                                              new_supported_args)
+
         try:
-            if hasattr(self._hook_instance.__class__, "settings"):
-                # this hook has a 'settings' property defined. it is expecting
-                # 'settings' to be passed to the processing method.
-                return self._hook_instance.process_current_session(
-                    self.settings, item)
-            else:
-                # the hook hasn't been updated to handle collector settings.
-                # call the method without a settings argument
-                return self._hook_instance.process_current_session(item)
+            return self._hook_instance.process_current_session(*process_args)
         except Exception:
             error_msg = traceback.format_exc()
             logger.error(

tests/fixtures/config/hooks/api_collector.py

@@ -19,7 +19,7 @@ class BasicSceneCollector(HookBaseClass):
     A basic collector that handles files and general objects.
     """
 
-    def process_current_session(self, settings, parent_item):
+    def process_current_session(self, settings, parent_item, collection_args):
         """
         Analyzes the current scene open in a DCC and parents a subtree of items
         under the parent_item passed in.
@@ -45,3 +45,6 @@ def process_current_session(self, settings, parent_item):
         )
 
         parent_item.local_properties.collector_property = "collector_property"
+
+        if collection_args:
+            parent_item.properties.update(collection_args)

tests/fixtures/config/hooks/collector.py

@@ -19,7 +19,7 @@ class BasicSceneCollector(HookBaseClass):
     A basic collector that handles files and general objects.
     """
 
-    def process_current_session(self, settings, parent_item):
+    def process_current_session(self, settings, parent_item, collection_args):
         """
         Analyzes the current scene open in a DCC and parents a subtree of items
         under the parent_item passed in.

tests/test_manager.py

@@ -154,3 +154,35 @@ def test_nodes():
 
         with self.assertRaisesRegex(Exception, "Test error!"):
             self.manager.publish(test_nodes())
+
+    def test_collection_args(self):
+        """
+        Ensures that collection_args is propagated to the collector hook
+        """
+        collection_args = {
+            "flag_1": True,
+            "flag_2": False,
+            "flag_3": 10,
+        }
+        # if collection_args is provided to collect_session
+        # api_collector adds collection_args to root item properties
+        # this is done only for test purposes
+
+        # without collection_args provided,
+        # root_item properties do not contain collection_args
+        self.manager.collect_session()
+        root_item = self.manager.tree.root_item
+        root_item_properties = root_item.properties
+        self.assertNotIn("flag_1", root_item_properties)
+        self.assertNotIn("flag_2", root_item_properties)
+        self.assertNotIn("flag_3", root_item_properties)
+
+        # now we provide collection_args
+        self.manager.collect_session(collection_args)
+
+        # so properties should be filled
+        self.assertEqual(root_item.properties["flag_1"], collection_args["flag_1"])
+        self.assertEqual(root_item.properties["flag_2"], collection_args["flag_2"])
+        self.assertEqual(root_item.properties["flag_3"], collection_args["flag_3"])
+
+