Fix bpy.data.user_map() ignoring unused datablocks.
[blender.git] / source / blender / python / intern / bpy_rna_id_collection.c
1 /*
2  * ***** BEGIN GPL LICENSE BLOCK *****
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License
6  * as published by the Free Software Foundation; either version 2
7  * of the License, or (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software Foundation,
16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17  *
18  * Contributor(s): Bastien Montagne
19  *
20  * ***** END GPL LICENSE BLOCK *****
21  */
22
23 /** \file blender/python/intern/bpy_rna_id_collection.c
24  *  \ingroup pythonintern
25  *
26  * This file adds some helpers related to ID/Main handling, that cannot fit well in RNA itself.
27  */
28
29 #include <Python.h>
30 #include <stddef.h>
31
32 #include "MEM_guardedalloc.h"
33
34 #include "BLI_utildefines.h"
35 #include "BLI_bitmap.h"
36
37 #include "BKE_global.h"
38 #include "BKE_main.h"
39 #include "BKE_library.h"
40 #include "BKE_library_query.h"
41
42 #include "DNA_ID.h"
43
44 #include "bpy_util.h"
45 #include "bpy_rna_id_collection.h"
46
47 #include "../generic/py_capi_utils.h"
48 #include "../generic/python_utildefines.h"
49
50 #include "RNA_access.h"
51 #include "RNA_types.h"
52 #include "RNA_enum_types.h"
53
54 #include "bpy_rna.h"
55
56 typedef struct IDUserMapData {
57         /* place-holder key only used for lookups to avoid creating new data only for lookups
58          * (never return its contents) */
59         PyObject *py_id_key_lookup_only;
60
61         /* we loop over data-blocks that this ID points to (do build a reverse lookup table) */
62         PyObject *py_id_curr;
63         ID *id_curr;
64
65         /* filter the values we add into the set */
66         BLI_bitmap *types_bitmap;
67
68         PyObject *user_map;     /* set to fill in as we iterate */
69         bool is_subset;         /* true when we're only mapping a subset of all the ID's (subset arg is passed) */
70 } IDUserMapData;
71
72
73 static int id_code_as_index(const short idcode)
74 {
75         return (int)*((unsigned short *)&idcode);
76 }
77
78 static bool id_check_type(const ID *id, const BLI_bitmap *types_bitmap)
79 {
80         return BLI_BITMAP_TEST_BOOL(types_bitmap, id_code_as_index(GS(id->name)));
81 }
82
83 static int foreach_libblock_id_user_map_callback(
84         void *user_data, ID *UNUSED(self_id), ID **id_p, int UNUSED(cb_flag))
85 {
86         IDUserMapData *data = user_data;
87
88         if (*id_p) {
89
90                 if (data->types_bitmap) {
91                         if (!id_check_type(*id_p, data->types_bitmap)) {
92                                 return IDWALK_RET_NOP;
93                         }
94                 }
95
96                 /* pyrna_struct_hash() uses ptr.data only,
97                  * but pyrna_struct_richcmp() uses also ptr.type,
98                  * so we need to create a valid PointerRNA here...
99                  */
100                 PyObject *key = data->py_id_key_lookup_only;
101                 RNA_id_pointer_create(*id_p, &((BPy_StructRNA *)key)->ptr);
102
103                 PyObject *set;
104                 if ((set = PyDict_GetItem(data->user_map, key)) == NULL) {
105
106                         /* limit to key's added already */
107                         if (data->is_subset) {
108                                 return IDWALK_RET_NOP;
109                         }
110
111                         /* Cannot use our placeholder key here! */
112                         key = pyrna_id_CreatePyObject(*id_p);
113                         set = PySet_New(NULL);
114                         PyDict_SetItem(data->user_map, key, set);
115                         Py_DECREF(set);
116                         Py_DECREF(key);
117                 }
118
119                 if (data->py_id_curr == NULL) {
120                         data->py_id_curr = pyrna_id_CreatePyObject(data->id_curr);
121                 }
122
123                 PySet_Add(set, data->py_id_curr);
124         }
125
126         return IDWALK_RET_NOP;
127 }
128
129 PyDoc_STRVAR(bpy_user_map_doc,
130 ".. method:: user_map([subset=(id1, id2, ...)], key_types={..}, value_types={..})\n"
131 "\n"
132 "   Returns a mapping of all ID datablocks in current ``bpy.data`` to a set of all datablocks using them.\n"
133 "\n"
134 "   For list of valid set members for key_types & value_types, see: :class:`bpy.types.KeyingSetPath.id_type`.\n"
135 "\n"
136 "   :arg subset: When passed, only these data-blocks and their users will be included as keys/values in the map.\n"
137 "   :type subset: sequence\n"
138 "   :arg key_types: Filter the keys mapped by ID types.\n"
139 "   :type key_types: set of strings\n"
140 "   :arg value_types: Filter the values in the set by ID types.\n"
141 "   :type value_types: set of strings\n"
142 "   :return: dictionary of :class:`bpy.types.ID` instances, with sets of ID's as their values.\n"
143 "   :rtype: dict\n"
144 );
145 static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *kwds)
146 {
147 #if 0  /* If someone knows how to get a proper 'self' in that case... */
148         BPy_StructRNA *pyrna = (BPy_StructRNA *)self;
149         Main *bmain = pyrna->ptr.data;
150 #else
151         Main *bmain = G.main;  /* XXX Ugly, but should work! */
152 #endif
153
154         static const char *kwlist[] = {"subset", "key_types", "value_types", NULL};
155         PyObject *subset = NULL;
156
157         PyObject *key_types = NULL;
158         PyObject *val_types = NULL;
159         BLI_bitmap *key_types_bitmap = NULL;
160         BLI_bitmap *val_types_bitmap = NULL;
161
162         PyObject *ret = NULL;
163
164
165         if (!PyArg_ParseTupleAndKeywords(
166                 args, kwds, "|O$O!O!:user_map", (char **)kwlist,
167                 &subset,
168                 &PySet_Type, &key_types,
169                 &PySet_Type, &val_types))
170         {
171                 return NULL;
172         }
173
174         if (key_types) {
175                 key_types_bitmap = pyrna_set_to_enum_bitmap(
176                         rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types");
177                 if (key_types_bitmap == NULL) {
178                         goto error;
179                 }
180         }
181
182         if (val_types) {
183                 val_types_bitmap = pyrna_set_to_enum_bitmap(
184                         rna_enum_id_type_items, val_types, sizeof(short), true, USHRT_MAX, "value types");
185                 if (val_types_bitmap == NULL) {
186                         goto error;
187                 }
188         }
189
190         IDUserMapData data_cb = {NULL};
191
192         if (subset) {
193                 PyObject *subset_fast = PySequence_Fast(subset, "user_map");
194                 if (subset_fast == NULL) {
195                         goto error;
196                 }
197
198                 PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast);
199                 Py_ssize_t subset_len = PySequence_Fast_GET_SIZE(subset_fast);
200
201                 data_cb.user_map = _PyDict_NewPresized(subset_len);
202                 data_cb.is_subset = true;
203                 for (; subset_len; subset_array++, subset_len--) {
204                         PyObject *set = PySet_New(NULL);
205                         PyDict_SetItem(data_cb.user_map, *subset_array, set);
206                         Py_DECREF(set);
207                 }
208                 Py_DECREF(subset_fast);
209         }
210         else {
211                 data_cb.user_map = PyDict_New();
212         }
213
214         data_cb.types_bitmap = key_types_bitmap;
215
216         ListBase *lb_array[MAX_LIBARRAY];
217         int lb_index;
218         lb_index = set_listbasepointers(bmain, lb_array);
219
220         while (lb_index--) {
221
222                 if (val_types_bitmap && lb_array[lb_index]->first) {
223                         if (!id_check_type(lb_array[lb_index]->first, val_types_bitmap)) {
224                                 continue;
225                         }
226                 }
227
228                 for (ID *id = lb_array[lb_index]->first; id; id = id->next) {
229                         /* One-time init, ID is just used as placeholder here, we abuse this in iterator callback
230                          * to avoid having to rebuild a complete bpyrna object each time for the key searching
231                          * (where only ID pointer value is used). */
232                         if (data_cb.py_id_key_lookup_only == NULL) {
233                                 data_cb.py_id_key_lookup_only = pyrna_id_CreatePyObject(id);
234                         }
235
236                         if (!data_cb.is_subset) {
237                                 PyObject *key = data_cb.py_id_key_lookup_only;
238                                 PyObject *set;
239
240                                 RNA_id_pointer_create(id, &((BPy_StructRNA *)key)->ptr);
241
242                                 /* We have to insert the key now, otherwise ID unused would be missing from final dict... */
243                                 if ((set = PyDict_GetItem(data_cb.user_map, key)) == NULL) {
244                                         /* Cannot use our placeholder key here! */
245                                         key = pyrna_id_CreatePyObject(id);
246                                         set = PySet_New(NULL);
247                                         PyDict_SetItem(data_cb.user_map, key, set);
248                                         Py_DECREF(set);
249                                         Py_DECREF(key);
250                                 }
251                         }
252
253                         data_cb.id_curr = id;
254                         BKE_library_foreach_ID_link(id, foreach_libblock_id_user_map_callback, &data_cb, IDWALK_NOP);
255
256                         if (data_cb.py_id_curr) {
257                                 Py_DECREF(data_cb.py_id_curr);
258                                 data_cb.py_id_curr = NULL;
259                         }
260                 }
261         }
262
263         ret = data_cb.user_map;
264
265
266 error:
267
268         Py_XDECREF(data_cb.py_id_key_lookup_only);
269
270         if (key_types_bitmap) {
271                 MEM_freeN(key_types_bitmap);
272         }
273
274         if (val_types_bitmap) {
275                 MEM_freeN(val_types_bitmap);
276         }
277
278         return ret;
279
280 }
281
282 int BPY_rna_id_collection_module(PyObject *mod_par)
283 {
284         static PyMethodDef user_map = {
285             "user_map", (PyCFunction)bpy_user_map, METH_VARARGS | METH_KEYWORDS, bpy_user_map_doc};
286
287         PyModule_AddObject(mod_par, "_rna_id_collection_user_map", PyCFunction_New(&user_map, NULL));
288
289         return 0;
290 }