diff --git a/source/blender/python/intern/bpy_rna_id_collection.cc b/source/blender/python/intern/bpy_rna_id_collection.cc index 3d5368c3b61..389730d78be 100644 --- a/source/blender/python/intern/bpy_rna_id_collection.cc +++ b/source/blender/python/intern/bpy_rna_id_collection.cc @@ -15,6 +15,7 @@ #include "BLI_bitmap.h" +#include "BKE_bpath.hh" #include "BKE_global.hh" #include "BKE_lib_id.hh" #include "BKE_lib_query.hh" @@ -282,6 +283,209 @@ error: return ret; } +struct IDFilePathMapData { + /* Data unchanged for the whole process. */ + + /** Set to fill in as we iterate. */ + PyObject *file_path_map; + + /** Whether to include library filepath of linked IDs or not. */ + bool include_libraries; + + /* Data modified for each processed ID. */ + + /** The processed ID. */ + ID *id; + /** The set of file paths for the processed ID. */ + PyObject *id_file_path_set; +}; + +static bool foreach_id_file_path_map_callback(BPathForeachPathData *bpath_data, + char * /*path_dst*/, + size_t /*path_dst_maxncpy*/, + const char *path_src) +{ + IDFilePathMapData &data = *static_cast(bpath_data->user_data); + ID *id = data.id; + PyObject *id_file_path_set = data.id_file_path_set; + + BLI_assert(id == bpath_data->owner_id); + + if (path_src && *path_src) { + PyObject *path = PyUnicode_FromString(path_src); + PySet_Add(id_file_path_set, path); + Py_DECREF(path); + } + return false; +} + +static void foreach_id_file_path_map(BPathForeachPathData &bpath_data) +{ + IDFilePathMapData &data = *static_cast(bpath_data.user_data); + ID *id = data.id; + PyObject *id_file_path_set = data.id_file_path_set; + + if (data.include_libraries && ID_IS_LINKED(id)) { + PyObject *path = PyUnicode_FromString(id->lib->filepath); + PySet_Add(id_file_path_set, path); + Py_DECREF(path); + } + + BKE_bpath_foreach_path_id(&bpath_data, id); +} + +PyDoc_STRVAR( + /* Wrap. */ + bpy_file_path_map_doc, + ".. method:: file_path_map(subset=None, key_types=None, include_libraries=False)\n" + "\n" + " Returns a mapping of all ID data-blocks in current ``bpy.data`` to a set of all " + "file paths used by them.\n" + "\n" + " For list of valid set members for key_types, see: " + ":class:`bpy.types.KeyingSetPath.id_type`.\n" + "\n" + " :arg subset: When given, only these data-blocks and their used file paths " + "will be included as keys/values in the map.\n" + " :type subset: sequence\n" + " :arg key_types: When given, filter the keys mapped by ID types. Ignored if ``subset`` is " + "also given.\n" + " :type key_types: set of strings\n" + " :arg include_libraries: Include library file paths of linked data. False by default.\n" + " :type include_libraries: bool\n" + " :return: dictionary of :class:`bpy.types.ID` instances, with sets of file path " + "strings as their values.\n" + " :rtype: dict\n"); +static PyObject *bpy_file_path_map(PyObject * /*self*/, PyObject *args, PyObject *kwds) +{ +#if 0 /* If someone knows how to get a proper 'self' in that case... */ + BPy_StructRNA *pyrna = (BPy_StructRNA *)self; + Main *bmain = pyrna->ptr.data; +#else + Main *bmain = G_MAIN; /* XXX Ugly, but should work! */ +#endif + + PyObject *subset = nullptr; + + PyObject *key_types = nullptr; + PyObject *include_libraries = nullptr; + BLI_bitmap *key_types_bitmap = nullptr; + + PyObject *ret = nullptr; + + IDFilePathMapData filepathmap_data{}; + BPathForeachPathData bpath_data{}; + + static const char *_keywords[] = {"subset", "key_types", "include_libraries", nullptr}; + static _PyArg_Parser _parser = { + PY_ARG_PARSER_HEAD_COMPAT() + "|$" /* Optional keyword only arguments. */ + "O" /* `subset` */ + "O!" /* `key_types` */ + "O!" /* `include_libraries` */ + ":file_path_map", + _keywords, + nullptr, + }; + if (!_PyArg_ParseTupleAndKeywordsFast(args, + kwds, + &_parser, + &subset, + &PySet_Type, + &key_types, + &PyBool_Type, + &include_libraries)) + { + return nullptr; + } + + if (key_types) { + key_types_bitmap = pyrna_enum_bitmap_from_set( + rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types"); + if (key_types_bitmap == nullptr) { + goto error; + } + } + + bpath_data.bmain = bmain; + bpath_data.callback_function = foreach_id_file_path_map_callback; + /* TODO: needs to be controlable from caller (add more options to the API). */ + bpath_data.flag = BKE_BPATH_FOREACH_PATH_SKIP_PACKED | BKE_BPATH_TRAVERSE_SKIP_WEAK_REFERENCES; + bpath_data.user_data = &filepathmap_data; + + filepathmap_data.include_libraries = (include_libraries == Py_True); + + if (subset) { + PyObject *subset_fast = PySequence_Fast(subset, "user_map"); + if (subset_fast == nullptr) { + goto error; + } + + PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast); + Py_ssize_t subset_len = PySequence_Fast_GET_SIZE(subset_fast); + + filepathmap_data.file_path_map = _PyDict_NewPresized(subset_len); + for (; subset_len; subset_array++, subset_len--) { + if (PyDict_Contains(filepathmap_data.file_path_map, *subset_array)) { + continue; + } + + ID *id; + if (!pyrna_id_FromPyObject(*subset_array, &id)) { + PyErr_Format(PyExc_TypeError, + "Expected an ID type in `subset` iterable, not %.200s", + Py_TYPE(*subset_array)->tp_name); + Py_DECREF(subset_fast); + Py_DECREF(filepathmap_data.file_path_map); + goto error; + } + + filepathmap_data.id_file_path_set = PySet_New(nullptr); + PyDict_SetItem( + filepathmap_data.file_path_map, *subset_array, filepathmap_data.id_file_path_set); + Py_DECREF(filepathmap_data.id_file_path_set); + + filepathmap_data.id = id; + foreach_id_file_path_map(bpath_data); + } + Py_DECREF(subset_fast); + } + else { + ListBase *lb; + ID *id; + filepathmap_data.file_path_map = PyDict_New(); + + FOREACH_MAIN_LISTBASE_BEGIN (bmain, lb) { + FOREACH_MAIN_LISTBASE_ID_BEGIN (lb, id) { + /* We can skip here in case we have some filter on key types. */ + if (key_types_bitmap && !id_check_type(id, key_types_bitmap)) { + break; + } + + PyObject *key = pyrna_id_CreatePyObject(id); + filepathmap_data.id_file_path_set = PySet_New(nullptr); + PyDict_SetItem(filepathmap_data.file_path_map, key, filepathmap_data.id_file_path_set); + Py_DECREF(filepathmap_data.id_file_path_set); + Py_DECREF(key); + + filepathmap_data.id = id; + foreach_id_file_path_map(bpath_data); + } + FOREACH_MAIN_LISTBASE_ID_END; + } + FOREACH_MAIN_LISTBASE_ID_END; + } + + ret = filepathmap_data.file_path_map; + +error: + if (key_types_bitmap != nullptr) { + MEM_freeN(key_types_bitmap); + } + + return ret; +} + PyDoc_STRVAR( /* Wrap. */ bpy_batch_remove_doc, @@ -440,6 +644,12 @@ PyMethodDef BPY_rna_id_collection_user_map_method_def = { METH_STATIC | METH_VARARGS | METH_KEYWORDS, bpy_user_map_doc, }; +PyMethodDef BPY_rna_id_collection_file_path_map_method_def = { + "file_path_map", + (PyCFunction)bpy_file_path_map, + METH_STATIC | METH_VARARGS | METH_KEYWORDS, + bpy_file_path_map_doc, +}; PyMethodDef BPY_rna_id_collection_batch_remove_method_def = { "batch_remove", (PyCFunction)bpy_batch_remove, diff --git a/source/blender/python/intern/bpy_rna_id_collection.hh b/source/blender/python/intern/bpy_rna_id_collection.hh index 2afcac8b87d..160555570f6 100644 --- a/source/blender/python/intern/bpy_rna_id_collection.hh +++ b/source/blender/python/intern/bpy_rna_id_collection.hh @@ -11,5 +11,6 @@ #include extern PyMethodDef BPY_rna_id_collection_user_map_method_def; +extern PyMethodDef BPY_rna_id_collection_file_path_map_method_def; extern PyMethodDef BPY_rna_id_collection_batch_remove_method_def; extern PyMethodDef BPY_rna_id_collection_orphans_purge_method_def; diff --git a/source/blender/python/intern/bpy_rna_types_capi.cc b/source/blender/python/intern/bpy_rna_types_capi.cc index 517f463336b..e04deec0ab9 100644 --- a/source/blender/python/intern/bpy_rna_types_capi.cc +++ b/source/blender/python/intern/bpy_rna_types_capi.cc @@ -45,6 +45,7 @@ static PyMethodDef pyrna_blenddata_methods[] = { {nullptr, nullptr, 0, nullptr}, /* #BPY_rna_id_collection_user_map_method_def */ + {nullptr, nullptr, 0, nullptr}, /* #BPY_rna_id_collection_file_path_map_method_def */ {nullptr, nullptr, 0, nullptr}, /* #BPY_rna_id_collection_batch_remove_method_def */ {nullptr, nullptr, 0, nullptr}, /* #BPY_rna_id_collection_orphans_purge_method_def */ {nullptr, nullptr, 0, nullptr}, /* #BPY_rna_data_context_method_def */ @@ -276,10 +277,11 @@ void BPY_rna_types_extend_capi() /* BlendData */ ARRAY_SET_ITEMS(pyrna_blenddata_methods, BPY_rna_id_collection_user_map_method_def, + BPY_rna_id_collection_file_path_map_method_def, BPY_rna_id_collection_batch_remove_method_def, BPY_rna_id_collection_orphans_purge_method_def, BPY_rna_data_context_method_def); - BLI_STATIC_ASSERT(ARRAY_SIZE(pyrna_blenddata_methods) == 5, "Unexpected number of methods") + BLI_STATIC_ASSERT(ARRAY_SIZE(pyrna_blenddata_methods) == 6, "Unexpected number of methods") pyrna_struct_type_extend_capi(&RNA_BlendData, pyrna_blenddata_methods, nullptr); /* BlendDataLibraries */ diff --git a/tests/python/bl_blendfile_relationships.py b/tests/python/bl_blendfile_relationships.py index e69be00b50f..5de8787ab3e 100644 --- a/tests/python/bl_blendfile_relationships.py +++ b/tests/python/bl_blendfile_relationships.py @@ -56,7 +56,7 @@ class TestBlendUserMap(TestBlendLibLinkHelper): } for k, v in expected_map.items(): self.assertIn(k, user_map) - self.assertEqual(user_map[k], v) + self.assertEqual(user_map[k], v, msg=f"ID {k.name} has unexpected user map") user_map = bpy.data.user_map(subset=[bpy.data.objects[0], bpy.data.meshes[0]]) expected_map = { @@ -66,11 +66,11 @@ class TestBlendUserMap(TestBlendLibLinkHelper): } for k, v in expected_map.items(): self.assertIn(k, user_map) - self.assertEqual(user_map[k], v) + self.assertEqual(user_map[k], v, msg=f"ID {k.name} has unexpected user map") user_map = bpy.data.user_map(key_types={'OBJECT', 'MESH'}) for k, v in expected_map.items(): self.assertIn(k, user_map) - self.assertEqual(user_map[k], v) + self.assertEqual(user_map[k], v, msg=f"ID {k.name} has unexpected user map") user_map = bpy.data.user_map(value_types={'SCENE'}) expected_map = { @@ -79,15 +79,96 @@ class TestBlendUserMap(TestBlendLibLinkHelper): } for k, v in expected_map.items(): self.assertIn(k, user_map) - self.assertEqual(user_map[k], v) + self.assertEqual(user_map[k], v, msg=f"ID {k.name} has unexpected user map") # Test handling of invalid parameters self.assertRaises(ValueError, bpy.data.user_map, value_types={'FOOBAR'}) self.assertRaises(TypeError, bpy.data.user_map, subset=[bpy.data.objects[0], bpy.data.meshes[0], "FooBar"]) +class TestBlendFilePathMap(TestBlendLibLinkHelper): + + def __init__(self, args): + super().__init__(args) + + def test_file_path_map(self): + def abspaths(file_path_map): + return {k: {os.path.normpath(bpy.path.abspath(p)) for p in v} + for k, v in file_path_map.items()} + + output_dir = self.args.output_dir + output_blendfile_path = self.init_lib_data_indirect_lib() + + # Simple link of a single ObData. + self.reset_blender() + + bpy.ops.wm.open_mainfile(filepath=output_blendfile_path) + + self.assertEqual(len(bpy.data.images), 1) + self.assertIsNotNone(bpy.data.images[0].library) + self.assertEqual(len(bpy.data.materials), 1) + self.assertIsNotNone(bpy.data.materials[0].library) + self.assertEqual(len(bpy.data.meshes), 1) + self.assertEqual(len(bpy.data.objects), 1) + self.assertEqual(len(bpy.data.collections), 1) + + blendlib_path = os.path.normpath(bpy.path.abspath(bpy.data.materials[0].library.filepath)) + image_path = os.path.join(native_pathsep(self.args.src_test_dir), + native_pathsep('imbuf_io/reference/jpeg-rgb-90__from__rgba08.jpg')) + + file_path_map = abspaths(bpy.data.file_path_map()) + # Note: Workspaces and screens are ignored here. + expected_map = { + bpy.data.images[0]: {image_path}, + bpy.data.materials[0]: set(), + bpy.data.scenes[0]: set(), + bpy.data.collections[0]: set(), + bpy.data.libraries[0]: {blendlib_path}, + bpy.data.meshes[0]: set(), + bpy.data.objects[0]: set(), + bpy.data.window_managers[0]: set(), + } + for k, v in expected_map.items(): + self.assertIn(k, file_path_map) + self.assertEqual(file_path_map[k], v, msg=f"ID {k.name} has unexpected filepath map") + + file_path_map = abspaths(bpy.data.file_path_map(include_libraries=True)) + # Note: Workspaces and screens are ignored here. + expected_map = { + bpy.data.images[0]: {image_path, blendlib_path}, + bpy.data.materials[0]: {blendlib_path}, + bpy.data.scenes[0]: set(), + bpy.data.collections[0]: set(), + bpy.data.libraries[0]: {blendlib_path}, + bpy.data.meshes[0]: set(), + bpy.data.objects[0]: set(), + bpy.data.window_managers[0]: set(), + } + for k, v in expected_map.items(): + self.assertIn(k, file_path_map) + self.assertEqual(file_path_map[k], v, msg=f"ID {k.name} has unexpected filepath map") + + file_path_map = abspaths(bpy.data.file_path_map(subset=[bpy.data.images[0], bpy.data.materials[0]])) + expected_map = { + bpy.data.images[0]: {image_path}, + bpy.data.materials[0]: set(), + } + for k, v in expected_map.items(): + self.assertIn(k, file_path_map) + self.assertEqual(file_path_map[k], v, msg=f"ID {k.name} has unexpected filepath map") + partial_map = abspaths(bpy.data.file_path_map(key_types={'IMAGE', 'MATERIAL'})) + for k, v in expected_map.items(): + self.assertIn(k, file_path_map) + self.assertEqual(file_path_map[k], v, msg=f"ID {k.name} has unexpected filepath map") + + # Test handling of invalid parameters + self.assertRaises(ValueError, bpy.data.file_path_map, key_types={'FOOBAR'}) + self.assertRaises(TypeError, bpy.data.file_path_map, subset=[bpy.data.objects[0], bpy.data.images[0], "FooBar"]) + + TESTS = ( TestBlendUserMap, + TestBlendFilePathMap, )