Update io_descriptor initialize on first use

python/tank/descriptor/io_descriptor/base.py

@@ -10,9 +10,11 @@
 
 import os
 import contextlib
+import traceback
 
 from .. import constants
 from ... import LogManager
+from ... import hook
 from ...util import filesystem, sgre as re
 from ...util.version import is_version_newer
 from ..errors import TankDescriptorError, TankMissingManifestError
@@ -37,6 +39,8 @@ class IODescriptorBase(object):
     Tank App store and one which knows how to handle the local file system.
     """
 
+    HOOK_NAME = "register_descriptor.py"
+    IO_DESCRIPTORS_INITIALIZED = False
     _factory = {}
 
     @classmethod
@@ -65,6 +69,10 @@ def create(cls, bundle_type, descriptor_dict, sg_connection):
         :returns: Instance of class deriving from :class:`IODescriptorBase`
         :raises: TankDescriptorError
         """
+        if not cls.IO_DESCRIPTORS_INITIALIZED:
+            cls.initialize_io_descriptor_types(
+                descriptor_dict, sg_connection, bundle_type
+            )
         descriptor_type = descriptor_dict.get("type")
         if descriptor_type not in cls._factory:
             raise TankDescriptorError(
@@ -73,6 +81,51 @@ def create(cls, bundle_type, descriptor_dict, sg_connection):
         class_obj = cls._factory[descriptor_type]
         return class_obj(descriptor_dict, sg_connection, bundle_type)
 
+    @classmethod
+    def initialize_io_descriptor_types(
+        cls, descriptor_dict, sg_connection, bundle_type
+    ):
+        # First put our base hook implementation into the array.
+        base_class_path = os.path.normpath(
+            os.path.join(
+                os.path.dirname(__file__),
+                os.pardir,
+                os.pardir,
+                os.pardir,
+                os.pardir,
+                "hooks",
+                cls.HOOK_NAME,
+            )
+        )
+        hook_inheritance_chain = [base_class_path]
+
+        # Then, check if there is a config-level override.
+        hook_path = os.path.join(
+            descriptor_dict.get("path"), "core", "hooks", cls.HOOK_NAME
+        )
+        if os.path.isfile(hook_path):
+            hook_inheritance_chain.append(hook_path)
+
+        try:
+            instance = hook.create_hook_instance(hook_inheritance_chain, parent=None)
+            instance.init(sg_connection, bundle_type, descriptor_dict)
+            instance.register_io_descriptors()
+        except TankDescriptorError:
+            from tank.descriptor.io_descriptor import _initialize_descriptor_factory
+
+            log.warning(
+                "Error while executing {hook_name} from {hook_path}. "
+                "Falling back to core descriptors".format(
+                    hook_name=cls.HOOK_NAME, hook_path=hook_inheritance_chain
+                )
+            )
+            log.debug(traceback.format_exc())
+            _initialize_descriptor_factory(cls)
+        finally:
+            # We only want to call this at most once.
+            # If it fails the first time, what could possible make it succeed a second time?
+            cls.IO_DESCRIPTORS_INITIALIZED = True
+
     def __init__(self, descriptor_dict, sg_connection, bundle_type):
         """
         Constructor

tests/descriptor_tests/test_io_descriptors.py

@@ -217,6 +217,40 @@ def test_cache_locations(self):
             ],
         )
 
+    def test_custom_io_descriptor_type(self):
+        config_root = os.path.join(self.fixtures_root, "config")
+        sg = self.mockgun
+
+        core = sgtk.descriptor.create_descriptor(
+            self.mockgun,
+            sgtk.descriptor.Descriptor.CONFIG,
+            {
+                "type": "dev",
+                "path": config_root,
+                "version": "v0.1.0",
+                "name": "tk-core",
+            },
+        )
+
+        location = {
+            "type": "bitbucket_release",
+            "organization": "tk-core-testing",
+            "repository": "tk-multi-demo",
+            "version": "v0.1.0",
+        }
+
+        downloads_root = os.path.join(self.fixtures_root, "config")
+        tk_bundle_app = sgtk.descriptor.create_descriptor(
+            sg,
+            sgtk.descriptor.Descriptor.APP,
+            location,
+            bundle_cache_root_override=downloads_root,
+        )
+        tk_bundle_app._io_descriptor.ensure_local()
+        self.assertTrue(
+            os.path.exists(tk_bundle_app._io_descriptor._get_primary_cache_path())
+        )
+
     def test_download_receipt(self):
         """
         Tests the download receipt logic

tests/fixtures/config/core/hooks/register_descriptor.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2020 Shotgun Software Inc.
+#
+# CONFIDENTIAL AND PROPRIETARY
+#
+# This work is provided "AS IS" and subject to the Shotgun Pipeline Toolkit
+# Source Code License included in this distribution package. See LICENSE.
+# By accessing, using, copying or modifying this work you indicate your
+# agreement to the Shotgun Pipeline Toolkit Source Code License. All rights
+# not expressly granted therein are reserved by Shotgun Software Inc.
+
+"""
+This hook is used override some of the functionality of the :class:`~sgtk.bootstrap.ToolkitManager`.
+
+It will be instantiated only after a configuration has been selected by the :class:`~sgtk.bootstrap.ToolkitManager`.
+Therefore, this hook will not be invoked to download a configuration. However, the Toolkit Core,
+applications, frameworks and engines can be downloaded through the hook.
+"""
+import os
+
+from sgtk import get_hook_baseclass
+from sgtk.descriptor.io_descriptor.downloadable import IODescriptorDownloadable
+from sgtk.descriptor.errors import TankError, TankDescriptorError
+from sgtk.util.shotgun import download
+
+
+class IODescriptorBitbucketRelease(IODescriptorDownloadable):
+    def __init__(self, descriptor_dict, sg_connection, bundle_type):
+        """
+        Constructor
+
+        :param descriptor_dict: descriptor dictionary describing the bundle
+        :param sg_connection: Shotgun connection to associated site.
+        :param bundle_type: Either AppDescriptor.APP, CORE, ENGINE or FRAMEWORK.
+        :return: Descriptor instance
+        """
+        super(IODescriptorBitbucketRelease, self).__init__(
+            descriptor_dict, sg_connection, bundle_type
+        )
+        self._validate_descriptor(
+            descriptor_dict,
+            required=["type", "organization", "repository", "version"],
+            optional=[],
+        )
+        self._sg_connection = sg_connection
+        self._bundle_type = bundle_type
+        self._organization = descriptor_dict["organization"]
+        self._repository = descriptor_dict["repository"]
+        self._version = descriptor_dict["version"]
+
+    def _get_bundle_cache_path(self, bundle_cache_root):
+        """
+        Given a cache root, compute a cache path suitable
+        for this descriptor, using the 0.18+ path format.
+
+        :param bundle_cache_root: Bundle cache root path
+        :return: Path to bundle cache location
+        """
+        return os.path.join(
+            bundle_cache_root,
+            "bitbucket",
+            self._organization,
+            self.get_system_name(),
+            self.get_version(),
+        )
+
+    def get_system_name(self):
+        """
+        Returns a short name, suitable for use in configuration files
+        and for folders on disk, e.g. 'tk-maya'
+        """
+        return self._repository
+
+    def get_version(self):
+        """
+        Returns the version number string for this item.
+        In this case, this is the linked shotgun attachment id.
+        """
+        return self._version
+
+    def _download_local(self, destination_path):
+        """
+        Retrieves this version to local repo.
+        Will exit early if app already exists local.
+
+        :param destination_path: The directory path to which the shotgun entity is to be
+        downloaded to.
+        """
+        url = "https://bitbucket.org/{organization}/{system_name}/get/{version}.zip"
+        url = url.format(
+            organization=self._organization,
+            system_name=self.get_system_name(),
+            version=self.get_version(),
+        )
+
+        try:
+            download.download_and_unpack_url(
+                self._sg_connection, url, destination_path, auto_detect_bundle=True
+            )
+        except TankError as e:
+            raise TankDescriptorError(
+                "Failed to download %s from %s. Error: %s" % (self, url, e)
+            )
+
+
+class MyCustomRegisterDescriptorsHook(get_hook_baseclass()):
+    def register_io_descriptors(self):
+        # To register the default IODescriptor classes
+        super(MyCustomRegisterDescriptorsHook, self).register_io_descriptors()
+        self.io_descriptor_base.register_descriptor_factory(
+            "bitbucket_release", IODescriptorBitbucketRelease
+        )