diff --git a/client/ayon_core/plugins/load/create_hero_version.py b/client/ayon_core/plugins/load/create_hero_version.py new file mode 100644 index 0000000000..7e1a0d8a3d --- /dev/null +++ b/client/ayon_core/plugins/load/create_hero_version.py @@ -0,0 +1,415 @@ + +import os +import copy +import shutil +import errno +import itertools +from concurrent.futures import ThreadPoolExecutor + +from speedcopy import copyfile +import clique +import ayon_api +from ayon_api.operations import OperationsSession, new_version_entity +from ayon_api.utils import create_entity_id +from qtpy import QtWidgets, QtCore +from ayon_core import style +from ayon_core.pipeline import load, Anatomy +from ayon_core.lib import create_hard_link, source_hash +from ayon_core.lib.file_transaction import wait_for_future_errors +from ayon_core.pipeline.publish import get_publish_template_name +from ayon_core.pipeline.template_data import get_template_data + + +def prepare_changes(old_entity, new_entity): + changes = {} + for key in set(new_entity.keys()): + if key == "attrib": + continue + if key in new_entity and new_entity[key] != old_entity.get(key): + changes[key] = new_entity[key] + attrib_changes = {} + if "attrib" in new_entity: + for key, value in new_entity["attrib"].items(): + if value != old_entity["attrib"].get(key): + attrib_changes[key] = value + if attrib_changes: + changes["attrib"] = attrib_changes + return changes + + +class CreateHeroVersion(load.ProductLoaderPlugin): + """Create hero version from selected context.""" + + is_multiple_contexts_compatible = False + representations = {"*"} + product_types = {"*"} + label = "Create Hero Version" + order = 36 + icon = "star" + color = "#ffd700" + + ignored_representation_names = [] + db_representation_context_keys = [ + "project", "folder", "asset", "hierarchy", "task", "product", + "subset", "family", "representation", "username", "user", "output" + ] + use_hardlinks = False + + def message(self, text): + msgBox = QtWidgets.QMessageBox() + msgBox.setText(text) + msgBox.setStyleSheet(style.load_stylesheet()) + msgBox.setWindowFlags( + msgBox.windowFlags() | QtCore.Qt.FramelessWindowHint + ) + msgBox.exec_() + + + def load(self, context, name=None, namespace=None, options=None) -> None: + """Load hero version from context (dict as in context.py).""" + success = True + errors = [] + + # Extract project, product, version, folder from context + project = context.get("project") + product = context.get("product") + version = context.get("version") + folder = context.get("folder") + task_entity = ayon_api.get_task_by_id( + task_id=version.get("taskId"), project_name=project["name"]) + + anatomy = Anatomy(project["name"]) + + version_id = version["id"] + project_name = project["name"] + repres = list( + ayon_api.get_representations(project_name, version_ids={version_id})) + anatomy_data = get_template_data( + project_entity=project, + folder_entity=folder, + task_entity=task_entity, + ) + anatomy_data["product"] = { + "name": product["name"], + "type": product["productType"], + } + published_representations = {} + for repre in repres: + repre_anatomy = anatomy_data + repre_anatomy["ext"] = repre.get("ext", "") + published_representations[repre["id"]] = { + "representation": repre, + "published_files": [f["path"] for f in repre.get("files", [])], + "anatomy_data": repre_anatomy + } + + instance_data = { + "productName": product["name"], + "productType": product["productType"], + "anatomyData": anatomy_data, + "publishDir": "", # TODO: Set to actual publish directory + "published_representations": published_representations, + "versionEntity": version, + } + + try: + self.create_hero_version(instance_data, anatomy, context) + except Exception as exc: + success = False + errors.append(str(exc)) + if success: + self.message("Hero version created successfully.") + else: + self.message( + f"Failed to create hero version:\n{chr(10).join(errors)}") + + def create_hero_version(self, instance_data, anatomy, context): + """Create hero version from instance data.""" + published_repres = instance_data.get("published_representations") + if not published_repres: + raise RuntimeError("No published representations found.") + + project_name = anatomy.project_name + template_key = get_publish_template_name( + project_name, + context.get("hostName"), + instance_data.get("productType"), + instance_data.get("anatomyData", {}).get("task", {}).get("name"), + instance_data.get("anatomyData", {}).get("task", {}).get("type"), + project_settings=context.get("project_settings", {}), + hero=True, + logger=None + ) + hero_template = anatomy.get_template_item("hero", template_key, "path", default=None) + if hero_template is None: + raise RuntimeError(f"Project anatomy does not have hero template key: {template_key}") + + print(f"Hero template: {hero_template.template}") + + hero_publish_dir = self.get_publish_dir(instance_data, anatomy, template_key) + + print(f"Hero publish dir: {hero_publish_dir}") + + src_version_entity = instance_data.get("versionEntity") + filtered_repre_ids = [] + for repre_id, repre_info in published_repres.items(): + repre = repre_info["representation"] + if repre["name"].lower() in self.ignored_representation_names: + filtered_repre_ids.append(repre_id) + for repre_id in filtered_repre_ids: + published_repres.pop(repre_id, None) + if not published_repres: + raise RuntimeError("All published representations were filtered by name.") + + if src_version_entity is None: + src_version_entity = self.version_from_representations(project_name, published_repres) + if not src_version_entity: + raise RuntimeError("Can't find origin version in database.") + if src_version_entity["version"] == 0: + raise RuntimeError("Version 0 cannot have hero version.") + + all_copied_files = [] + transfers = instance_data.get("transfers", list()) + for _src, dst in transfers: + dst = os.path.normpath(dst) + if dst not in all_copied_files: + all_copied_files.append(dst) + hardlinks = instance_data.get("hardlinks", list()) + for _src, dst in hardlinks: + dst = os.path.normpath(dst) + if dst not in all_copied_files: + all_copied_files.append(dst) + + all_repre_file_paths = [] + for repre_info in published_repres.values(): + published_files = repre_info.get("published_files") or [] + for file_path in published_files: + file_path = os.path.normpath(file_path) + if file_path not in all_repre_file_paths: + all_repre_file_paths.append(file_path) + + instance_publish_dir = os.path.normpath(instance_data["publishDir"]) + other_file_paths_mapping = [] + for file_path in all_copied_files: + if not file_path.startswith(instance_publish_dir): + continue + if file_path in all_repre_file_paths: + continue + dst_filepath = file_path.replace(instance_publish_dir, hero_publish_dir) + other_file_paths_mapping.append((file_path, dst_filepath)) + + old_version, old_repres = self.current_hero_ents(project_name, src_version_entity) + inactive_old_repres_by_name = {} + old_repres_by_name = {} + for repre in old_repres: + low_name = repre["name"].lower() + if repre["active"]: + old_repres_by_name[low_name] = repre + else: + inactive_old_repres_by_name[low_name] = repre + + op_session = OperationsSession() + entity_id = old_version["id"] if old_version else None + new_hero_version = new_version_entity( + -src_version_entity["version"], + src_version_entity["productId"], + task_id=src_version_entity.get("taskId"), + data=copy.deepcopy(src_version_entity["data"]), + attribs=copy.deepcopy(src_version_entity["attrib"]), + entity_id=entity_id, + ) + if old_version: + update_data = prepare_changes(old_version, new_hero_version) + op_session.update_entity(project_name, "version", old_version["id"], update_data) + else: + op_session.create_entity(project_name, "version", new_hero_version) + + # Store hero entity to instance_data + instance_data["heroVersionEntity"] = new_hero_version + + old_repres_to_replace = {} + old_repres_to_delete = {} + for repre_info in published_repres.values(): + repre = repre_info["representation"] + repre_name_low = repre["name"].lower() + if repre_name_low in old_repres_by_name: + old_repres_to_replace[repre_name_low] = old_repres_by_name.pop(repre_name_low) + if old_repres_by_name: + old_repres_to_delete = old_repres_by_name + + backup_hero_publish_dir = None + if os.path.exists(hero_publish_dir): + backup_hero_publish_dir = hero_publish_dir + ".BACKUP" + max_idx = 10 + idx = 0 + _backup_hero_publish_dir = backup_hero_publish_dir + while os.path.exists(_backup_hero_publish_dir): + try: + shutil.rmtree(_backup_hero_publish_dir) + backup_hero_publish_dir = _backup_hero_publish_dir + break + except Exception: + _backup_hero_publish_dir = backup_hero_publish_dir + str(idx) + if not os.path.exists(_backup_hero_publish_dir): + backup_hero_publish_dir = _backup_hero_publish_dir + break + if idx > max_idx: + raise AssertionError(f"Backup folders are fully occupied to max index {max_idx}") + idx += 1 + try: + os.rename(hero_publish_dir, backup_hero_publish_dir) + except PermissionError: + raise AssertionError( + "Could not create hero version because it is " + "not possible to replace current hero files.") + + try: + src_to_dst_file_paths = [] + repre_integrate_data = [] + path_template_obj = anatomy.get_template_item( + "hero", template_key, "path") + for repre_info in published_repres.values(): + published_files = repre_info["published_files"] + if len(published_files) == 0: + continue + anatomy_data = copy.deepcopy(repre_info["anatomy_data"]) + anatomy_data.pop("version", None) + template_filled = path_template_obj.format_strict(anatomy_data) + repre_context = template_filled.used_values + for key in self.db_representation_context_keys: + value = anatomy_data.get(key) + if value is not None: + repre_context[key] = value + repre_entity = copy.deepcopy(repre_info["representation"]) + repre_entity.pop("id", None) + repre_entity["versionId"] = new_hero_version["id"] + repre_entity["context"] = repre_context + repre_entity["attrib"] = { + "path": str(template_filled), + "template": hero_template.template + } + dst_paths = [] + if len(published_files) == 1: + dst_paths.append(str(template_filled)) + src_to_dst_file_paths.append((published_files[0], template_filled)) + print(f"Single published file: {published_files[0]} -> {template_filled}") + else: + collections, remainders = clique.assemble(published_files) + if remainders or not collections or len(collections) > 1: + raise Exception( + "Integrity error. Files of published representation is " + "combination of frame collections and single files.") + src_col = collections[0] + frame_splitter = "_-_FRAME_SPLIT_-_" + anatomy_data["frame"] = frame_splitter + _template_filled = path_template_obj.format_strict(anatomy_data) + head, tail = _template_filled.split(frame_splitter) + padding = anatomy.templates_obj.frame_padding + dst_col = clique.Collection(head=head, padding=padding, tail=tail) + dst_col.indexes.clear() + dst_col.indexes.update(src_col.indexes) + for src_file, dst_file in zip(src_col, dst_col): + src_to_dst_file_paths.append((src_file, dst_file)) + dst_paths.append(dst_file) + print(f"Collection published file: {src_file} -> {dst_file}") + repre_integrate_data.append((repre_entity, dst_paths)) + + # Copy files + with ThreadPoolExecutor(max_workers=8) as executor: + futures = [ + executor.submit(self.copy_file, src_path, dst_path) + for src_path, dst_path in itertools.chain(src_to_dst_file_paths, other_file_paths_mapping) + ] + wait_for_future_errors(executor, futures) + + # Update/create representations + for repre_entity, dst_paths in repre_integrate_data: + repre_files = self.get_files_info(dst_paths, anatomy) + repre_entity["files"] = repre_files + repre_name_low = repre_entity["name"].lower() + if repre_name_low in old_repres_to_replace: + old_repre = old_repres_to_replace.pop(repre_name_low) + repre_entity["id"] = old_repre["id"] + update_data = prepare_changes(old_repre, repre_entity) + op_session.update_entity(project_name, "representation", old_repre["id"], update_data) + elif repre_name_low in inactive_old_repres_by_name: + inactive_repre = inactive_old_repres_by_name.pop(repre_name_low) + repre_entity["id"] = inactive_repre["id"] + update_data = prepare_changes(inactive_repre, repre_entity) + op_session.update_entity(project_name, "representation", inactive_repre["id"], update_data) + else: + op_session.create_entity(project_name, "representation", repre_entity) + + for repre in old_repres_to_delete.values(): + op_session.update_entity(project_name, "representation", repre["id"], {"active": False}) + + op_session.commit() + + if backup_hero_publish_dir is not None and os.path.exists(backup_hero_publish_dir): + shutil.rmtree(backup_hero_publish_dir) + + except Exception: + if backup_hero_publish_dir is not None and os.path.exists(backup_hero_publish_dir): + if os.path.exists(hero_publish_dir): + shutil.rmtree(hero_publish_dir) + os.rename(backup_hero_publish_dir, hero_publish_dir) + raise + + def get_files_info(self, filepaths, anatomy): + file_infos = [] + for filepath in filepaths: + file_info = self.prepare_file_info(filepath, anatomy) + file_infos.append(file_info) + return file_infos + + def prepare_file_info(self, path, anatomy): + return { + "id": create_entity_id(), + "name": os.path.basename(path), + "path": self.get_rootless_path(anatomy, path), + "size": os.path.getsize(path), + "hash": source_hash(path), + "hash_type": "op3", + } + + def get_publish_dir(self, instance_data, anatomy, template_key): + template_data = copy.deepcopy(instance_data.get("anatomyData", {})) + if "originalBasename" in instance_data: + template_data["originalBasename"] = instance_data["originalBasename"] + template_obj = anatomy.get_template_item("hero", template_key, "directory") + return os.path.normpath(template_obj.format_strict(template_data)) + + def get_rootless_path(self, anatomy, path): + success, rootless_path = anatomy.find_root_template_from_path(path) + if success: + path = rootless_path + return path + + def copy_file(self, src_path, dst_path): + dirname = os.path.dirname(dst_path) + try: + os.makedirs(dirname) + except OSError as exc: + if exc.errno != errno.EEXIST: + raise + if self.use_hardlinks: + try: + create_hard_link(src_path, dst_path) + return + except OSError as exc: + if exc.errno not in [errno.EXDEV, errno.EINVAL]: + raise + copyfile(src_path, dst_path) + + def version_from_representations(self, project_name, repres): + for repre_info in repres.values(): + version = ayon_api.get_version_by_id(project_name, repre_info["representation"]["versionId"]) + if version: + return version + + def current_hero_ents(self, project_name, version): + hero_version = ayon_api.get_hero_version_by_product_id(project_name, version["productId"]) + if not hero_version: + return (None, []) + hero_repres = list(ayon_api.get_representations(project_name, version_ids={hero_version["id"]})) + return (hero_version, hero_repres)