a7bd5afd4225c1d20ea84907573d6a6680f611ca
[blender-staging.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 /* Those folowing are only to support hack of not listing some internal 'backward' pointers in generated user_map... */
44 #include "DNA_object_types.h"
45 #include "DNA_key_types.h"
46
47 #include "bpy_util.h"
48 #include "bpy_rna_id_collection.h"
49
50 #include "../generic/py_capi_utils.h"
51 #include "../generic/python_utildefines.h"
52
53 #include "RNA_access.h"
54 #include "RNA_types.h"
55 #include "RNA_enum_types.h"
56
57 #include "bpy_rna.h"
58
59 typedef struct IDUserMapData {
60         /* place-holder key only used for lookups to avoid creating new data only for lookups
61          * (never return its contents) */
62         PyObject *py_id_key_lookup_only;
63
64         /* we loop over data-blocks that this ID points to (do build a reverse lookup table) */
65         PyObject *py_id_curr;
66         ID *id_curr;
67
68         /* filter the values we add into the set */
69         BLI_bitmap *types_bitmap;
70
71         PyObject *user_map;     /* set to fill in as we iterate */
72         bool is_subset;         /* true when we're only mapping a subset of all the ID's (subset arg is passed) */
73 } IDUserMapData;
74
75
76 static int id_code_as_index(const short idcode)
77 {
78         return (int)*((unsigned short *)&idcode);
79 }
80
81 static bool id_check_type(const ID *id, const BLI_bitmap *types_bitmap)
82 {
83         return BLI_BITMAP_TEST_BOOL(types_bitmap, id_code_as_index(GS(id->name)));
84 }
85
86 static int foreach_libblock_id_user_map_callback(
87         void *user_data, ID *self_id, ID **id_p, int UNUSED(cb_flag))
88 {
89         IDUserMapData *data = user_data;
90
91         if (*id_p) {
92
93                 if (data->types_bitmap) {
94                         if (!id_check_type(*id_p, data->types_bitmap)) {
95                                 return IDWALK_RET_NOP;
96                         }
97                 }
98
99                 if ((GS(self_id->name) == ID_OB) && (id_p == (ID **)&((Object *)self_id)->proxy_from)) {
100                         /* We skip proxy_from here, since it some internal pointer which is not irrelevant info for py/API level. */
101                         return IDWALK_RET_NOP;
102                 }
103                 else if ((GS(self_id->name) == ID_KE) && (id_p == (ID **)&((Key *)self_id)->from)) {
104                         /* We skip from here, since it some internal pointer which is not irrelevant info for py/API level. */
105                         return IDWALK_RET_NOP;
106                 }
107
108                 /* pyrna_struct_hash() uses ptr.data only,
109                  * but pyrna_struct_richcmp() uses also ptr.type,
110                  * so we need to create a valid PointerRNA here...
111                  */
112                 PyObject *key = data->py_id_key_lookup_only;
113                 RNA_id_pointer_create(*id_p, &((BPy_StructRNA *)key)->ptr);
114
115                 PyObject *set;
116                 if ((set = PyDict_GetItem(data->user_map, key)) == NULL) {
117
118                         /* limit to key's added already */
119                         if (data->is_subset) {
120                                 return IDWALK_RET_NOP;
121                         }
122
123                         /* Cannot use our placeholder key here! */
124                         key = pyrna_id_CreatePyObject(*id_p);
125                         set = PySet_New(NULL);
126                         PyDict_SetItem(data->user_map, key, set);
127                         Py_DECREF(set);
128                         Py_DECREF(key);
129                 }
130
131                 if (data->py_id_curr == NULL) {
132                         data->py_id_curr = pyrna_id_CreatePyObject(data->id_curr);
133                 }
134
135                 PySet_Add(set, data->py_id_curr);
136         }
137
138         return IDWALK_RET_NOP;
139 }
140
141 PyDoc_STRVAR(bpy_user_map_doc,
142 ".. method:: user_map([subset=(id1, id2, ...)], key_types={..}, value_types={..})\n"
143 "\n"
144 "   Returns a mapping of all ID datablocks in current ``bpy.data`` to a set of all datablocks using them.\n"
145 "\n"
146 "   For list of valid set members for key_types & value_types, see: :class:`bpy.types.KeyingSetPath.id_type`.\n"
147 "\n"
148 "   :arg subset: When passed, only these data-blocks and their users will be included as keys/values in the map.\n"
149 "   :type subset: sequence\n"
150 "   :arg key_types: Filter the keys mapped by ID types.\n"
151 "   :type key_types: set of strings\n"
152 "   :arg value_types: Filter the values in the set by ID types.\n"
153 "   :type value_types: set of strings\n"
154 "   :return: dictionary of :class:`bpy.types.ID` instances, with sets of ID's as their values.\n"
155 "   :rtype: dict\n"
156 );
157 static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *kwds)
158 {
159 #if 0  /* If someone knows how to get a proper 'self' in that case... */
160         BPy_StructRNA *pyrna = (BPy_StructRNA *)self;
161         Main *bmain = pyrna->ptr.data;
162 #else
163         Main *bmain = G.main;  /* XXX Ugly, but should work! */
164 #endif
165
166         static const char *kwlist[] = {"subset", "key_types", "value_types", NULL};
167         PyObject *subset = NULL;
168
169         PyObject *key_types = NULL;
170         PyObject *val_types = NULL;
171         BLI_bitmap *key_types_bitmap = NULL;
172         BLI_bitmap *val_types_bitmap = NULL;
173
174         PyObject *ret = NULL;
175
176
177         if (!PyArg_ParseTupleAndKeywords(
178                 args, kwds, "|O$O!O!:user_map", (char **)kwlist,
179                 &subset,
180                 &PySet_Type, &key_types,
181                 &PySet_Type, &val_types))
182         {
183                 return NULL;
184         }
185
186         if (key_types) {
187                 key_types_bitmap = pyrna_set_to_enum_bitmap(
188                         rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types");
189                 if (key_types_bitmap == NULL) {
190                         goto error;
191                 }
192         }
193
194         if (val_types) {
195                 val_types_bitmap = pyrna_set_to_enum_bitmap(
196                         rna_enum_id_type_items, val_types, sizeof(short), true, USHRT_MAX, "value types");
197                 if (val_types_bitmap == NULL) {
198                         goto error;
199                 }
200         }
201
202         IDUserMapData data_cb = {NULL};
203
204         if (subset) {
205                 PyObject *subset_fast = PySequence_Fast(subset, "user_map");
206                 if (subset_fast == NULL) {
207                         goto error;
208                 }
209
210                 PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast);
211                 Py_ssize_t subset_len = PySequence_Fast_GET_SIZE(subset_fast);
212
213                 data_cb.user_map = _PyDict_NewPresized(subset_len);
214                 data_cb.is_subset = true;
215                 for (; subset_len; subset_array++, subset_len--) {
216                         PyObject *set = PySet_New(NULL);
217                         PyDict_SetItem(data_cb.user_map, *subset_array, set);
218                         Py_DECREF(set);
219                 }
220                 Py_DECREF(subset_fast);
221         }
222         else {
223                 data_cb.user_map = PyDict_New();
224         }
225
226         data_cb.types_bitmap = key_types_bitmap;
227
228         ListBase *lb_array[MAX_LIBARRAY];
229         int lb_index;
230         lb_index = set_listbasepointers(bmain, lb_array);
231
232         while (lb_index--) {
233
234                 if (val_types_bitmap && lb_array[lb_index]->first) {
235                         if (!id_check_type(lb_array[lb_index]->first, val_types_bitmap)) {
236                                 continue;
237                         }
238                 }
239
240                 for (ID *id = lb_array[lb_index]->first; id; id = id->next) {
241                         /* One-time init, ID is just used as placeholder here, we abuse this in iterator callback
242                          * to avoid having to rebuild a complete bpyrna object each time for the key searching
243                          * (where only ID pointer value is used). */
244                         if (data_cb.py_id_key_lookup_only == NULL) {
245                                 data_cb.py_id_key_lookup_only = pyrna_id_CreatePyObject(id);
246                         }
247
248                         if (!data_cb.is_subset) {
249                                 PyObject *key = data_cb.py_id_key_lookup_only;
250                                 PyObject *set;
251
252                                 RNA_id_pointer_create(id, &((BPy_StructRNA *)key)->ptr);
253
254                                 /* We have to insert the key now, otherwise ID unused would be missing from final dict... */
255                                 if ((set = PyDict_GetItem(data_cb.user_map, key)) == NULL) {
256                                         /* Cannot use our placeholder key here! */
257                                         key = pyrna_id_CreatePyObject(id);
258                                         set = PySet_New(NULL);
259                                         PyDict_SetItem(data_cb.user_map, key, set);
260                                         Py_DECREF(set);
261                                         Py_DECREF(key);
262                                 }
263                         }
264
265                         data_cb.id_curr = id;
266                         BKE_library_foreach_ID_link(NULL, id, foreach_libblock_id_user_map_callback, &data_cb, IDWALK_NOP);
267
268                         if (data_cb.py_id_curr) {
269                                 Py_DECREF(data_cb.py_id_curr);
270                                 data_cb.py_id_curr = NULL;
271                         }
272                 }
273         }
274
275         ret = data_cb.user_map;
276
277
278 error:
279
280         Py_XDECREF(data_cb.py_id_key_lookup_only);
281
282         if (key_types_bitmap) {
283                 MEM_freeN(key_types_bitmap);
284         }
285
286         if (val_types_bitmap) {
287                 MEM_freeN(val_types_bitmap);
288         }
289
290         return ret;
291
292 }
293
294 int BPY_rna_id_collection_module(PyObject *mod_par)
295 {
296         static PyMethodDef user_map = {
297             "user_map", (PyCFunction)bpy_user_map, METH_VARARGS | METH_KEYWORDS, bpy_user_map_doc};
298
299         PyModule_AddObject(mod_par, "_rna_id_collection_user_map", PyCFunction_New(&user_map, NULL));
300
301         return 0;
302 }