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