'users of ID' py API.
[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 bool foreach_libblock_id_user_map_callback(void *user_data, ID **id_p, int UNUSED(cb_flag))
84 {
85         IDUserMapData *data = user_data;
86
87         if (*id_p) {
88
89                 if (data->types_bitmap) {
90                         if (!id_check_type(*id_p, data->types_bitmap)) {
91                                 return true;
92                         }
93                 }
94
95                 /* pyrna_struct_hash() uses ptr.data only,
96                  * but pyrna_struct_richcmp() uses also ptr.type,
97                  * so we need to create a valid PointerRNA here...
98                  */
99                 PyObject *key = data->py_id_key_lookup_only;
100                 RNA_id_pointer_create(*id_p, &((BPy_StructRNA *)key)->ptr);
101
102                 PyObject *set;
103                 if ((set = PyDict_GetItem(data->user_map, key)) == NULL) {
104
105                         /* limit to key's added already */
106                         if (data->is_subset) {
107                                 return true;
108                         }
109
110                         /* Cannot use our placeholder key here! */
111                         key = pyrna_id_CreatePyObject(*id_p);
112                         set = PySet_New(NULL);
113                         PyDict_SetItem(data->user_map, key, set);
114                         Py_DECREF(set);
115                         Py_DECREF(key);
116                 }
117
118                 if (data->py_id_curr == NULL) {
119                         data->py_id_curr = pyrna_id_CreatePyObject(data->id_curr);
120                 }
121
122                 PySet_Add(set, data->py_id_curr);
123         }
124
125         return true;
126 }
127
128 PyDoc_STRVAR(bpy_user_map_doc,
129 ".. method:: user_map([subset=(id1, id2, ...)], key_types={..}, value_types={..})\n"
130 "\n"
131 "   Returns a mapping of all ID datablocks in current ``bpy.data`` to a set of all datablocks using them.\n"
132 "\n"
133 "   For list of valid set members for key_types & value_types, see: :class:`bpy.types.KeyingSetPath.id_type`.\n"
134 "\n"
135 "   :arg subset: When passed, only these data-blocks and their users will be included as keys/values in the map.\n"
136 "   :type subset: sequence\n"
137 "   :arg key_types: Filter the keys mapped by ID types.\n"
138 "   :type key_types: set of strings\n"
139 "   :arg value_types: Filter the values in the set by ID types.\n"
140 "   :type value_types: set of strings\n"
141 "   :return: dictionary of :class:`bpy.types.ID` instances, with sets of ID's as their values.\n"
142 "   :rtype: dict\n"
143 );
144 static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *kwds)
145 {
146 #if 0  /* If someone knows how to get a proper 'self' in that case... */
147         BPy_StructRNA *pyrna = (BPy_StructRNA *)self;
148         Main *bmain = pyrna->ptr.data;
149 #else
150         Main *bmain = G.main;  /* XXX Ugly, but should work! */
151 #endif
152
153         static const char *kwlist[] = {"subset", "key_types", "value_types", NULL};
154         PyObject *subset = NULL;
155
156         PyObject *key_types = NULL;
157         PyObject *val_types = NULL;
158         BLI_bitmap *key_types_bitmap = NULL;
159         BLI_bitmap *val_types_bitmap = NULL;
160
161         PyObject *ret = NULL;
162
163
164         if (!PyArg_ParseTupleAndKeywords(
165                 args, kwds, "|O$O!O!:user_map", (char **)kwlist,
166                 &subset,
167                 &PySet_Type, &key_types,
168                 &PySet_Type, &val_types))
169         {
170                 return NULL;
171         }
172
173         if (key_types) {
174                 key_types_bitmap = pyrna_set_to_enum_bitmap(
175                         rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types");
176                 if (key_types_bitmap == NULL) {
177                         goto error;
178                 }
179         }
180
181         if (val_types) {
182                 val_types_bitmap = pyrna_set_to_enum_bitmap(
183                         rna_enum_id_type_items, val_types, sizeof(short), true, USHRT_MAX, "value types");
184                 if (val_types_bitmap == NULL) {
185                         goto error;
186                 }
187         }
188
189         IDUserMapData data_cb = {NULL};
190
191         if (subset) {
192                 PyObject *subset_fast = PySequence_Fast(subset, "user_map");
193                 if (subset_fast == NULL) {
194                         goto error;
195                 }
196
197                 PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast);
198                 Py_ssize_t subset_len = PySequence_Fast_GET_SIZE(subset_fast);
199
200                 data_cb.user_map = PyDict_New();
201                 data_cb.is_subset = true;
202                 for (; subset_len; subset_array++, subset_len--) {
203                         PyObject *set = PySet_New(NULL);
204                         PyDict_SetItem(data_cb.user_map, *subset_array, set);
205                         Py_DECREF(set);
206                 }
207                 Py_DECREF(subset_fast);
208         }
209         else {
210                 data_cb.user_map = PyDict_New();
211         }
212
213         data_cb.types_bitmap = key_types_bitmap;
214
215         ListBase *lb_array[MAX_LIBARRAY];
216         int lb_index;
217         lb_index = set_listbasepointers(bmain, lb_array);
218
219         while (lb_index--) {
220
221                 if (val_types_bitmap && lb_array[lb_index]->first) {
222                         if (!id_check_type(lb_array[lb_index]->first, val_types_bitmap)) {
223                                 continue;
224                         }
225                 }
226
227                 for (ID *id = lb_array[lb_index]->first; id; id = id->next) {
228                         /* One-time init, ID is just used as placeholder here, we abuse this in iterator callback
229                          * to avoid having to rebuild a complete bpyrna object each time for the key searching
230                          * (where only ID pointer value is used). */
231                         if (data_cb.py_id_key_lookup_only == NULL) {
232                                 data_cb.py_id_key_lookup_only = pyrna_id_CreatePyObject(id);
233                         }
234
235                         data_cb.id_curr = id;
236                         BKE_library_foreach_ID_link(id, foreach_libblock_id_user_map_callback, &data_cb, IDWALK_NOP);
237                         if (data_cb.py_id_curr) {
238                                 Py_DECREF(data_cb.py_id_curr);
239                                 data_cb.py_id_curr = NULL;
240                         }
241                 }
242         }
243
244         ret = data_cb.user_map;
245
246
247 error:
248
249         Py_XDECREF(data_cb.py_id_key_lookup_only);
250
251         if (key_types_bitmap) {
252                 MEM_freeN(key_types_bitmap);
253         }
254
255         if (val_types_bitmap) {
256                 MEM_freeN(val_types_bitmap);
257         }
258
259         return ret;
260
261 }
262
263 int BPY_rna_id_collection_module(PyObject *mod_par)
264 {
265         static PyMethodDef user_map = {
266             "user_map", (PyCFunction)bpy_user_map, METH_VARARGS | METH_KEYWORDS, bpy_user_map_doc};
267
268         PyModule_AddObject(mod_par, "_rna_id_collection_user_map", PyCFunction_New(&user_map, NULL));
269
270         return 0;
271 }