OP-2951 - added function to collect referenced representation ids

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