bpy_extras: Add utilities for getting ID references

An alternate to D14839, implemented in Python and
relying on bpy.data.user_map(). That function
gives us a mapping of what ID is referenced by
what set of IDs. The inverse of this would also
be useful, which is now available from
bpy_extras.id_map_utils.get_id_reference_map().

From there, we can use get_all_referenced_ids()
to get a set of all IDs referenced by a given ID
either directly or indirectly.

To get only the direct references, we can simply
pass the ID of interest as a key to the dictionary
returned from get_id_reference_map().

Reviewed By: mont29

Differential Revision: https://developer.blender.org/D14843

2 files changed, 50 insertions, 0 deletions

diff --git a/release/scripts/modules/bpy_extras/__init__.py b/release/scripts/modules/bpy_extras/__init__.py
index 1af9048ebfd..15a8d00cddc 100644
--- a/release/scripts/modules/bpy_extras/__init__.py
+++ b/release/scripts/modules/bpy_extras/__init__.py
@@ -16,4 +16,5 @@ __all__ = (
     "mesh_utils",
     "node_utils",
     "view3d_utils",
+    "id_map_utils",
 )
diff --git a/release/scripts/modules/bpy_extras/id_map_utils.py b/release/scripts/modules/bpy_extras/id_map_utils.py
new file mode 100644
index 00000000000..cf39f2185c6
--- /dev/null
+++ b/release/scripts/modules/bpy_extras/id_map_utils.py
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+
+# <pep8 compliant>
+
+from typing import Dict, Set
+import bpy
+from bpy.types import ID
+
+
+__all__ = (
+    "get_id_reference_map",
+    "get_all_referenced_ids",
+)
+
+
+def get_id_reference_map() -> Dict[ID, Set[ID]]:
+    """Return a dictionary of direct datablock references for every datablock in the blend file."""
+    inv_map = {}
+    for key, values in bpy.data.user_map().items():
+        for value in values:
+            if value == key:
+                # So an object is not considered to be referencing itself.
+                continue
+            inv_map.setdefault(value, set()).add(key)
+    return inv_map
+
+
+def recursive_get_referenced_ids(
+    ref_map: Dict[ID, Set[ID]], id: ID, referenced_ids: Set, visited: Set
+):
+    """Recursively populate referenced_ids with IDs referenced by id."""
+    if id in visited:
+        # Avoid infinite recursion from circular references.
+        return
+    visited.add(id)
+    for ref in ref_map.get(id, []):
+        referenced_ids.add(ref)
+        recursive_get_referenced_ids(
+            ref_map=ref_map, id=ref, referenced_ids=referenced_ids, visited=visited
+        )
+
+
+def get_all_referenced_ids(id: ID, ref_map: Dict[ID, Set[ID]]) -> Set[ID]:
+    """Return a set of IDs directly or indirectly referenced by id."""
+    referenced_ids = set()
+    recursive_get_referenced_ids(
+        ref_map=ref_map, id=id, referenced_ids=referenced_ids, visited=set()
+    )
+    return referenced_ids