collect_last_version_repres implemented to get all representations of latest versions for all subsets of entered assets

pype/lib.py

@@ -670,3 +670,98 @@ def get_workfile_build_presets(task_name):
         break
 
     return per_task_preset
+
+
+def collect_last_version_repres(asset_entities):
+    """Collect subsets, versions and representations for asset_entities.
+
+    :param asset_entities: Asset entities for which want to find data
+    :type asset_entities: list
+    :return: collected entities
+    :rtype: dict
+
+    Example output:
+    ```
+    {
+        {Asset ID}: {
+            "asset_entity": <AssetEntity>,
+            "subsets": {
+                {Subset ID}: {
+                    "subset_entity": <SubsetEntity>,
+                    "version": {
+                        "version_entity": <VersionEntity>,
+                        "repres": [
+                            <RepreEntity1>, <RepreEntity2>, ...
+                        ]
+                    }
+                },
+                ...
+            }
+        },
+        ...
+    }
+    output[asset_id]["subsets"][subset_id]["version"]["repres"]
+    ```
+    """
+
+    if not asset_entities:
+        return {}
+
+    asset_entity_by_ids = {asset["_id"]: asset for asset in asset_entities}
+
+    subsets = list(io.find({
+        "type": "subset",
+        "parent": {"$in": asset_entity_by_ids.keys()}
+    }))
+    subset_entity_by_ids = {subset["_id"]: subset for subset in subsets}
+
+    sorted_versions = list(io.find({
+        "type": "version",
+        "parent": {"$in": subset_entity_by_ids.keys()}
+    }).sort("name", -1))
+
+    subset_id_with_latest_version = []
+    last_versions_by_id = {}
+    for version in sorted_versions:
+        subset_id = version["parent"]
+        if subset_id in subset_id_with_latest_version:
+            continue
+        subset_id_with_latest_version.append(subset_id)
+        last_versions_by_id[version["_id"]] = version
+
+    repres = io.find({
+        "type": "representation",
+        "parent": {"$in": last_versions_by_id.keys()}
+    })
+
+    output = {}
+    for repre in repres:
+        version_id = repre["parent"]
+        version = last_versions_by_id[version_id]
+
+        subset_id = version["parent"]
+        subset = subset_entity_by_ids[subset_id]
+
+        asset_id = subset["parent"]
+        asset = asset_entity_by_ids[asset_id]
+
+        if asset_id not in output:
+            output[asset_id] = {
+                "asset_entity": asset,
+                "subsets": {}
+            }
+
+        if subset_id not in output[asset_id]["subsets"]:
+            output[asset_id]["subsets"][subset_id] = {
+                "subset_entity": subset,
+                "version": {
+                    "version_entity": version,
+                    "repres": []
+                }
+            }
+
+        output[asset_id]["subsets"][subset_id]["version"]["repres"].append(
+            repre
+        )
+
+    return output