diff --git a/openpype/lib/avalon_context.py b/openpype/lib/avalon_context.py index b4e6abb72d..e8a365ec39 100644 --- a/openpype/lib/avalon_context.py +++ b/openpype/lib/avalon_context.py @@ -1971,3 +1971,123 @@ def get_last_workfile( return os.path.normpath(os.path.join(workdir, filename)) return filename + + +@with_avalon +def get_linked_ids_for_representations(project, repre_ids, dbcon=None, + link_type=None, max_depth=0): + """Returns list of linked ids of particular type (if provided). + + Goes from representations to version, back to representations + Args: + project (str) + repre_ids (list) or (ObjectId) + dbcon (avalon.mongodb.AvalonMongoDB, optional): Avalon Mongo connection + with Session. + link_type (str): ['reference', '..] + max_depth (int): limit how many levels of recursion + Returns: + (list) of ObjectId - linked representations + """ + if not dbcon: + log.debug("Using `avalon.io` for query.") + dbcon = avalon.io + # Make sure is installed + dbcon.install() + + if dbcon.Session["AVALON_PROJECT"] != project: + dbcon.Session["AVALON_PROJECT"] = project + + if not isinstance(repre_ids, list): + repre_ids = [repre_ids] + + versions = avalon.io.find( + { + "_id": {"$in": repre_ids}, + "type": "representation" + }, + projection={"parent": True} + ) + version_ids = [version["parent"] for version in versions] + + graph_lookup = { + "from": project, + "startWith": "$data.inputLinks.id", + "connectFromField": "data.inputLinks.id", + "connectToField": "_id", + "as": "outputs_recursive", + "depthField": "depth" + } + if max_depth != 0: + # We offset by -1 since 0 basically means no recursion + # but the recursion only happens after the initial lookup + # for outputs. + graph_lookup["maxDepth"] = max_depth - 1 + + match = { + "_id": {"$in": version_ids}, + "type": "version" + } + + pipeline_ = [ + # Match + {"$match": match}, + # Recursive graph lookup for inputs + {"$graphLookup": graph_lookup} + ] + + result = dbcon.aggregate(pipeline_) + referenced_version_ids = _process_referenced_pipeline_result(result, + link_type) + + representations = avalon.io.find( + { + "parent": {"$in": list(referenced_version_ids)}, + "type": "representation" + }, + projection={"_id": True} + ) + ref_ids = {representation["_id"] for representation in representations} + return list(ref_ids) + + +def _process_referenced_pipeline_result(result, link_type): + """Filters result from pipeline for particular link_type. + + Pipeline cannot use link_type directly in a query. + Returns: + (list) + """ + referenced_version_ids = set() + correctly_linked_ids = set() + for item in result: + correctly_linked_ids = _filter_input_links(item["data"]["inputLinks"], + link_type, + correctly_linked_ids) + + # outputs_recursive in random order, sort by _id + outputs_recursive = sorted(item.get("outputs_recursive", []), + key=lambda d: d["_id"]) + # go from oldest to newest + # only older _id can reference another newer _id + for output in outputs_recursive[::-1]: + if output["_id"] not in correctly_linked_ids: # leaf + continue + + correctly_linked_ids = _filter_input_links( + output["data"].get("inputLinks", []), + link_type, + correctly_linked_ids) + + referenced_version_ids.add(output["_id"]) + + return referenced_version_ids + + +def _filter_input_links(input_links, link_type, correctly_linked_ids): + for input_link in input_links: + if not link_type or input_link["type"] == link_type: + correctly_linked_ids.add(input_link.get("id") or + input_link.get("_id")) # legacy + + return correctly_linked_ids