Patch D133: Python wrapper for BLI_kdtree (adds mathutils.kdtree)
[blender.git] / source / blender / python / mathutils / mathutils_kdtree.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): Dan Eicher, Campbell Barton
19  *
20  * ***** END GPL LICENSE BLOCK *****
21  */
22
23 /** \file blender/python/mathutils/mathutils_kdtree.c
24  *  \ingroup mathutils
25  *
26  * This file defines the 'mathutils.kdtree' module, a general purpose module to access
27  * blenders kdtree for 3d spatial lookups.
28  */
29
30 #include <Python.h>
31
32 #include "MEM_guardedalloc.h"
33
34 #include "BLI_utildefines.h"
35 #include "BLI_kdtree.h"
36
37 #include "../generic/py_capi_utils.h"
38 #include "mathutils.h"
39
40 #include "BLI_strict_flags.h"
41
42 typedef struct {
43         PyObject_HEAD
44         KDTree *obj;
45         unsigned int maxsize;
46         unsigned int count;
47         unsigned int count_balance;  /* size when we last balanced */
48 } PyKDTree;
49
50
51 /* -------------------------------------------------------------------- */
52 /* Utility helper functions */
53
54 static void kdtree_nearest_to_py_tuple(const KDTreeNearest *nearest, PyObject *py_retval)
55 {
56         BLI_assert(nearest->index >= 0);
57         BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
58
59         PyTuple_SET_ITEM(py_retval, 0, Vector_CreatePyObject((float *)nearest->co, 3, Py_NEW, NULL));
60         PyTuple_SET_ITEM(py_retval, 1, PyLong_FromLong(nearest->index));
61         PyTuple_SET_ITEM(py_retval, 2, PyFloat_FromDouble(nearest->dist));
62 }
63
64 static PyObject *kdtree_nearest_to_py(const KDTreeNearest *nearest)
65 {
66         PyObject *py_retval;
67
68         py_retval = PyTuple_New(3);
69
70         kdtree_nearest_to_py_tuple(nearest, py_retval);
71
72         return py_retval;
73 }
74
75 static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest *nearest)
76 {
77         PyObject *py_retval;
78
79         py_retval = PyTuple_New(3);
80
81         if (nearest->index != -1) {
82                 kdtree_nearest_to_py_tuple(nearest, py_retval);
83         }
84         else {
85                 PyC_Tuple_Fill(py_retval, Py_None);
86         }
87
88         return py_retval;
89 }
90
91
92 /* -------------------------------------------------------------------- */
93 /* KDTree */
94
95 /* annoying since arg parsing won't check overflow */
96 #define UINT_IS_NEG(n) ((n) > INT_MAX)
97
98 static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
99 {
100         unsigned int maxsize;
101         const char *keywords[] = {"size", NULL};
102
103         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *)"I:KDTree", (char **)keywords, &maxsize)) {
104                 return -1;
105         }
106
107         if (UINT_IS_NEG(maxsize)) {
108                 PyErr_SetString(PyExc_ValueError, "negative 'size' given");
109                 return -1;
110         }
111
112         self->obj = BLI_kdtree_new(maxsize);
113         self->maxsize = maxsize;
114         self->count = 0;
115         self->count_balance = 0;
116
117         return 0;
118 }
119
120 static void PyKDTree__tp_dealloc(PyKDTree *self)
121 {
122         BLI_kdtree_free(self->obj);
123         Py_TYPE(self)->tp_free((PyObject *)self);
124 }
125
126 PyDoc_STRVAR(py_kdtree_insert_doc,
127 ".. method:: insert(index, co)\n"
128 "\n"
129 "   Insert a point into the KDTree.\n"
130 "\n"
131 "   :arg co: Point 3d position.\n"
132 "   :type co: float triplet\n"
133 "   :arg index: The index of the point.\n"
134 "   :type index: int\n"
135 );
136 static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
137 {
138         PyObject *py_co;
139         float co[3];
140         int index;
141         const char *keywords[] = {"co", "index", NULL};
142
143         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Oi:insert", (char **)keywords,
144                                          &py_co, &index))
145         {
146                 return NULL;
147         }
148
149         if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1)
150                 return NULL;
151
152         if (index < 0) {
153                 PyErr_SetString(PyExc_ValueError, "negative index given");
154                 return NULL;
155         }
156
157         if (self->count >= self->maxsize) {
158                 PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
159                 return NULL;
160         }
161
162         BLI_kdtree_insert(self->obj, index, co, NULL);
163         self->count++;
164
165         Py_RETURN_NONE;
166 }
167
168 PyDoc_STRVAR(py_kdtree_balance_doc,
169 ".. method:: balance()\n"
170 "\n"
171 "   Balance the tree.\n"
172 );
173 static PyObject *py_kdtree_balance(PyKDTree *self)
174 {
175         BLI_kdtree_balance(self->obj);
176         self->count_balance = self->count;
177         Py_RETURN_NONE;
178 }
179
180 PyDoc_STRVAR(py_kdtree_find_doc,
181 ".. method:: find(co)\n"
182 "\n"
183 "   Find nearest point to ``co``.\n"
184 "\n"
185 "   :arg co: 3d coordinates.\n"
186 "   :type co: float triplet\n"
187 "   :return: Returns (:class:`Vector`, index, distance).\n"
188 "   :rtype: :class:`tuple`\n"
189 );
190 static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
191 {
192         PyObject *py_co;
193         float co[3];
194         KDTreeNearest nearest;
195         const char *keywords[] = {"co", NULL};
196
197         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "O:find", (char **)keywords,
198                                          &py_co))
199         {
200                 return NULL;
201         }
202
203         if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1)
204                 return NULL;
205
206         if (self->count != self->count_balance) {
207                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
208                 return NULL;
209         }
210
211
212         nearest.index = -1;
213
214         BLI_kdtree_find_nearest(self->obj, co, NULL, &nearest);
215
216         return kdtree_nearest_to_py_and_check(&nearest);
217 }
218
219 PyDoc_STRVAR(py_kdtree_find_n_doc,
220 ".. method:: find_n(co, n)\n"
221 "\n"
222 "   Find nearest ``n`` points to ``co``.\n"
223 "\n"
224 "   :arg co: 3d coordinates.\n"
225 "   :type co: float triplet\n"
226 "   :arg n: Number of points to find.\n"
227 "   :type n: int\n"
228 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
229 "   :rtype: :class:`list`\n"
230 );
231 static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
232 {
233         PyObject *py_list;
234         PyObject *py_co;
235         float co[3];
236         KDTreeNearest *nearest;
237         unsigned int n;
238         int i, found;
239         const char *keywords[] = {"co", "n", NULL};
240
241         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "OI:find_n", (char **)keywords,
242                                          &py_co, &n))
243         {
244                 return NULL;
245         }
246
247         if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1)
248                 return NULL;
249
250         if (UINT_IS_NEG(n)) {
251                 PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
252                 return NULL;
253         }
254
255         if (self->count != self->count_balance) {
256                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
257                 return NULL;
258         }
259
260         nearest = MEM_mallocN(sizeof(KDTreeNearest) * n, __func__);
261
262         found = BLI_kdtree_find_nearest_n(self->obj, co, NULL, nearest, n);
263
264         py_list = PyList_New(found);
265
266         for (i = 0; i < found; i++) {
267                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
268         }
269
270         MEM_freeN(nearest);
271
272         return py_list;
273 }
274
275 PyDoc_STRVAR(py_kdtree_find_range_doc,
276 ".. method:: find_range(co, radius)\n"
277 "\n"
278 "   Find all points within ``radius`` of ``co``.\n"
279 "\n"
280 "   :arg co: 3d coordinates.\n"
281 "   :type co: float triplet\n"
282 "   :arg radius: Distance to search for points.\n"
283 "   :type radius: float\n"
284 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
285 "   :rtype: :class:`list`\n"
286 );
287 static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
288 {
289         PyObject *py_list;
290         PyObject *py_co;
291         float co[3];
292         KDTreeNearest *nearest = NULL;
293         float radius;
294         int i, found;
295
296         const char *keywords[] = {"co", "radius", NULL};
297
298         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Of:find_range", (char **)keywords,
299                                          &py_co, &radius))
300         {
301                 return NULL;
302         }
303
304         if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1)
305                 return NULL;
306
307         if (radius < 0.0f) {
308                 PyErr_SetString(PyExc_RuntimeError, "negative radius given");
309                 return NULL;
310         }
311
312         if (self->count != self->count_balance) {
313                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
314                 return NULL;
315         }
316
317         found = BLI_kdtree_range_search(self->obj, co, NULL, &nearest, radius);
318
319         py_list = PyList_New(found);
320
321         for (i = 0; i < found; i++) {
322                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
323         }
324
325         if (nearest) {
326                 MEM_freeN(nearest);
327         }
328
329         return py_list;
330 }
331
332
333 static PyMethodDef PyKDTree_methods[] = {
334         {"insert", (PyCFunction)py_kdtree_insert, METH_VARARGS | METH_KEYWORDS, py_kdtree_insert_doc},
335         {"balance", (PyCFunction)py_kdtree_balance, METH_NOARGS, py_kdtree_balance_doc},
336         {"find", (PyCFunction)py_kdtree_find, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_doc},
337         {"find_n", (PyCFunction)py_kdtree_find_n, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_n_doc},
338         {"find_range", (PyCFunction)py_kdtree_find_range, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_range_doc},
339         {NULL, NULL, 0, NULL}
340 };
341
342 PyDoc_STRVAR(py_KDtree_doc,
343 "KdTree(size) -> new kd-tree initialized to hold ``size`` items.\n"
344 "\n"
345 ".. note::\n"
346 "\n"
347 "   :class:`KDTree.balance` must have been called before using any of the ``find`` methods.\n"
348 );
349 PyTypeObject PyKDTree_Type = {
350         PyVarObject_HEAD_INIT(NULL, 0)
351         "KDTree",                                    /* tp_name */
352         sizeof(PyKDTree),                            /* tp_basicsize */
353         0,                                           /* tp_itemsize */
354         /* methods */
355         (destructor)PyKDTree__tp_dealloc,            /* tp_dealloc */
356         NULL,                                        /* tp_print */
357         NULL,                                        /* tp_getattr */
358         NULL,                                        /* tp_setattr */
359         NULL,                                        /* tp_compare */
360         NULL,                                        /* tp_repr */
361         NULL,                                        /* tp_as_number */
362         NULL,                                        /* tp_as_sequence */
363         NULL,                                        /* tp_as_mapping */
364         NULL,                                        /* tp_hash */
365         NULL,                                        /* tp_call */
366         NULL,                                        /* tp_str */
367         NULL,                                        /* tp_getattro */
368         NULL,                                        /* tp_setattro */
369         NULL,                                        /* tp_as_buffer */
370         Py_TPFLAGS_DEFAULT,                          /* tp_flags */
371         py_KDtree_doc,                               /* Documentation string */
372         NULL,                                        /* tp_traverse */
373         NULL,                                        /* tp_clear */
374         NULL,                                        /* tp_richcompare */
375         0,                                           /* tp_weaklistoffset */
376         NULL,                                        /* tp_iter */
377         NULL,                                        /* tp_iternext */
378         (struct PyMethodDef *)PyKDTree_methods,      /* tp_methods */
379         NULL,                                        /* tp_members */
380         NULL,                                        /* tp_getset */
381         NULL,                                        /* tp_base */
382         NULL,                                        /* tp_dict */
383         NULL,                                        /* tp_descr_get */
384         NULL,                                        /* tp_descr_set */
385         0,                                           /* tp_dictoffset */
386         (initproc)PyKDTree__tp_init,                 /* tp_init */
387         (allocfunc)PyType_GenericAlloc,              /* tp_alloc */
388         (newfunc)PyType_GenericNew,                  /* tp_new */
389         (freefunc)0,                                 /* tp_free */
390         NULL,                                        /* tp_is_gc */
391         NULL,                                        /* tp_bases */
392         NULL,                                        /* tp_mro */
393         NULL,                                        /* tp_cache */
394         NULL,                                        /* tp_subclasses */
395         NULL,                                        /* tp_weaklist */
396         (destructor) NULL                            /* tp_del */
397 };
398
399 PyDoc_STRVAR(py_kdtree_doc,
400 "Generic 3-dimentional kd-tree to perform spatial searches."
401 );
402 static struct PyModuleDef kdtree_moduledef = {
403         PyModuleDef_HEAD_INIT,
404         "mathutils.kdtree",                          /* m_name */
405         py_kdtree_doc,                               /* m_doc */
406         0,                                           /* m_size */
407         NULL,                                        /* m_methods */
408         NULL,                                        /* m_reload */
409         NULL,                                        /* m_traverse */
410         NULL,                                        /* m_clear */
411         NULL                                         /* m_free */
412 };
413
414 PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
415 {
416         PyObject *m = PyModule_Create(&kdtree_moduledef);
417
418         if (m == NULL) {
419                 return NULL;
420         }
421
422         /* Register the 'KDTree' class */
423         if (PyType_Ready(&PyKDTree_Type)) {
424                 return NULL;
425         }
426         PyModule_AddObject(m, (char *)"KDTree", (PyObject *) &PyKDTree_Type);
427
428         return m;
429 }