[Bf-blender-cvs] [ef3ac50] tmp-id-users: Add key and value type filters

Campbell Barton noreply at git.blender.org
Wed Jan 6 19:07:11 CET 2016


Commit: ef3ac5073e9b7aee22f6ea7cbfa74de3b68ab6c4
Author: Campbell Barton
Date:   Thu Jan 7 04:58:42 2016 +1100
Branches: tmp-id-users
https://developer.blender.org/rBef3ac5073e9b7aee22f6ea7cbfa74de3b68ab6c4

Add key and value type filters

allows for eg:

    bpy.data.user_map(key_types={'MESH'}, value_types={'OBJECT', 'MATERIAL'})

Without this, return values could get very large.

===================================================================

M	source/blender/python/intern/bpy_rna_id_collection.c

===================================================================

diff --git a/source/blender/python/intern/bpy_rna_id_collection.c b/source/blender/python/intern/bpy_rna_id_collection.c
index 16e2f38..745f55b 100644
--- a/source/blender/python/intern/bpy_rna_id_collection.c
+++ b/source/blender/python/intern/bpy_rna_id_collection.c
@@ -29,7 +29,10 @@
 #include <Python.h>
 #include <stddef.h>
 
+#include "MEM_guardedalloc.h"
+
 #include "BLI_utildefines.h"
+#include "BLI_bitmap.h"
 
 #include "BKE_global.h"
 #include "BKE_main.h"
@@ -46,6 +49,8 @@
 
 #include "RNA_access.h"
 #include "RNA_types.h"
+#include "RNA_enum_types.h"
+
 #include "bpy_rna.h"
 
 typedef struct IDUserMapData {
@@ -57,24 +62,44 @@ typedef struct IDUserMapData {
 	PyObject *py_id_curr;
 	ID *id_curr;
 
+	/* filter the values we add into the set */
+	BLI_bitmap *types_bitmap;
+
 	PyObject *user_map;     /* set to fill in as we iterate */
 	bool is_subset;         /* true when we're only mapping a subset of all the ID's (subset arg is passed) */
 } IDUserMapData;
 
+
+static int id_code_as_index(const short idcode)
+{
+	return (int)*((unsigned short *)&idcode);
+}
+
+static bool id_check_type(const ID *id, const BLI_bitmap *types_bitmap)
+{
+	return BLI_BITMAP_TEST_BOOL(types_bitmap, id_code_as_index(GS(id->name)));
+}
+
 static bool foreach_libblock_id_user_map_callback(void *user_data, ID **id_p, int UNUSED(cb_flag))
 {
 	IDUserMapData *data = user_data;
 
 	if (*id_p) {
-		PyObject *set;
-		PyObject *key = data->py_id_key_lookup_only;
+
+		if (data->types_bitmap) {
+			if (!id_check_type(*id_p, data->types_bitmap)) {
+				return true;
+			}
+		}
 
 		/* pyrna_struct_hash() uses ptr.data only,
 		 * but pyrna_struct_richcmp() uses also ptr.type,
 		 * so we need to create a valid PointerRNA here...
 		 */
+		PyObject *key = data->py_id_key_lookup_only;
 		RNA_id_pointer_create(*id_p, &((BPy_StructRNA *)key)->ptr);
 
+		PyObject *set;
 		if ((set = PyDict_GetItem(data->user_map, key)) == NULL) {
 
 			/* limit to key's added already */
@@ -100,13 +125,64 @@ static bool foreach_libblock_id_user_map_callback(void *user_data, ID **id_p, in
 	return true;
 }
 
+static BLI_bitmap *pyrna_set_to_bitmap(
+        EnumPropertyItem *items, PyObject *value,
+        int type_size, bool type_convert_sign,
+        int bitmap_size,
+        const char *error_prefix)
+{
+	/* set looping */
+	Py_ssize_t pos = 0;
+	Py_ssize_t hash = 0;
+	PyObject *key;
+
+	BLI_bitmap *bitmap = BLI_BITMAP_NEW(bitmap_size, __func__);
+
+	while (_PySet_NextEntry(value, &pos, &key, &hash)) {
+		const char *param = _PyUnicode_AsString(key);
+		if (param == NULL) {
+			PyErr_Format(PyExc_TypeError,
+			             "%.200s expected a string, not %.200s",
+			             error_prefix, Py_TYPE(key)->tp_name);
+			goto error;
+		}
+
+		int ret;
+		if (pyrna_enum_value_from_id(items, param, &ret, error_prefix) == -1) {
+			goto error;
+		}
+
+		int index = ret;
+		if (type_convert_sign) {
+			if (type_size == 2) {
+				union { signed short as_sign; unsigned short as_unsign; } ret_convert;
+				ret_convert.as_sign = (short)ret;
+				index = (int)ret_convert.as_unsign;
+			}
+		}
+		BLI_assert(index < bitmap_size);
+		BLI_BITMAP_ENABLE(bitmap, index);
+	}
+
+	return bitmap;
+
+error:
+	MEM_freeN(bitmap);
+	return NULL;
+}
+
+
 PyDoc_STRVAR(bpy_user_map_doc,
-".. method:: user_map([subset=(id1, id2, ...)])\n"
+".. method:: user_map([subset=(id1, id2, ...)], key_types={..}, value_types={..})\n"
 "\n"
 "   Returns a mapping of all ID datablocks in current ``bpy.data`` to a set of all datablocks using them.\n"
 "\n"
 "   :arg subset: When passed, only these data-blocks and their users will be included as keys/values in the map.\n"
 "   :type subset: sequence\n"
+"   :arg key_types: Filter the keys mapped by ID types.\n"
+"   :type key_types: set of strings\n"
+"   :arg key_types: Filter the values in the set by ID types.\n"
+"   :type key_types: set of strings\n"
 "   :return: dictionary of :class:`bpy.types.ID` instances, with sets of ID's as their values.\n"
 "   :rtype: dict\n"
 );
@@ -119,22 +195,56 @@ static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *
 	Main *bmain = G.main;  /* XXX Ugly, but should work! */
 #endif
 
-	static const char *kwlist[] = {"subset", NULL};
+	static const char *kwlist[] = {"subset", "key_types", "value_types", NULL};
 	PyObject *subset = NULL;
 
-	IDUserMapData data_cb = {NULL};
+	PyObject *key_types = NULL;
+	PyObject *val_types = NULL;
+	BLI_bitmap *key_types_bitmap = NULL;
+	BLI_bitmap *val_types_bitmap = NULL;
+
+	PyObject *ret = NULL;
+
 
 	if (!PyArg_ParseTupleAndKeywords(
-	        args, kwds, "|O:user_map", (char **)kwlist,
-	        &subset))
+	        args, kwds, "|O$O!O!:user_map", (char **)kwlist,
+	        &subset,
+	        &PySet_Type, &key_types,
+	        &PySet_Type, &val_types))
 	{
 		return NULL;
 	}
 
+	if (key_types) {
+		key_types_bitmap = pyrna_set_to_bitmap(
+		        rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "types");
+		if (key_types_bitmap == NULL) {
+			goto error;
+		}
+	}
+
+	if (key_types) {
+		key_types_bitmap = pyrna_set_to_bitmap(
+		        rna_enum_id_type_items, key_types, sizeof(short), true, USHRT_MAX, "key types");
+		if (key_types_bitmap == NULL) {
+			goto error;
+		}
+	}
+
+	if (val_types) {
+		val_types_bitmap = pyrna_set_to_bitmap(
+		        rna_enum_id_type_items, val_types, sizeof(short), true, USHRT_MAX, "value types");
+		if (val_types_bitmap == NULL) {
+			goto error;
+		}
+	}
+
+	IDUserMapData data_cb = {NULL};
+
 	if (subset) {
 		PyObject *subset_fast = PySequence_Fast(subset, "user_map");
 		if (subset_fast == NULL) {
-			return NULL;
+			goto error;
 		}
 
 		PyObject **subset_array = PySequence_Fast_ITEMS(subset_fast);
@@ -153,11 +263,20 @@ static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *
 		data_cb.user_map = PyDict_New();
 	}
 
+	data_cb.types_bitmap = val_types_bitmap;
+
 	ListBase *lb_array[MAX_LIBARRAY];
 	int lb_index;
 	lb_index = set_listbasepointers(bmain, lb_array);
 
 	while (lb_index--) {
+
+		if (key_types_bitmap && lb_array[lb_index]->first) {
+			if (!id_check_type(lb_array[lb_index]->first, key_types_bitmap)) {
+				continue;
+			}
+		}
+
 		for (ID *id = lb_array[lb_index]->first; id; id = id->next) {
 			/* One-time init, ID is just used as placeholder here, we abuse this in iterator callback
 			 * to avoid having to rebuild a complete bpyrna object each time for the key searching
@@ -175,9 +294,23 @@ static PyObject *bpy_user_map(PyObject *UNUSED(self), PyObject *args, PyObject *
 		}
 	}
 
+	ret = data_cb.user_map;
+
+
+error:
+
 	Py_XDECREF(data_cb.py_id_key_lookup_only);
 
-	return data_cb.user_map;
+	if (key_types_bitmap) {
+		MEM_freeN(key_types_bitmap);
+	}
+
+	if (val_types_bitmap) {
+		MEM_freeN(val_types_bitmap);
+	}
+
+	return ret;
+
 }
 
 int BPY_rna_id_collection_module(PyObject *mod_par)




More information about the Bf-blender-cvs mailing list