@@ -39,7 +39,6 @@ static size_t Size(const void *node, const char *prop) {
3939 PyObject * o = Attribute (node , prop );
4040 if (o != NULL ) {
4141 retval = PySequence_Size (o );
42- Py_DECREF (o );
4342 }
4443
4544 return retval ;
@@ -287,25 +286,42 @@ static void PyUastIter_dealloc(PyObject *self)
287286 UastIteratorFree (((PyUastIter * )self )-> iter );
288287}
289288
289+ static bool initFilter (PyObject * args , PyObject * * obj , const char * * query )
290+ {
291+ if (!PyArg_ParseTuple (args , "Os" , obj , query )) {
292+ return false;
293+ }
294+
295+ itemAtAllocsList = PyList_New (0 );
296+ stringAllocsList = PyList_New (0 );
297+ return true;
298+ }
299+
300+ static void cleanupFilter (void )
301+ {
302+ Py_DECREF (itemAtAllocsList );
303+ Py_DECREF (stringAllocsList );
304+ }
305+
306+ static void filterError (void )
307+ {
308+ char * error = LastError ();
309+ PyErr_SetString (PyExc_RuntimeError , error );
310+ free (error );
311+ cleanupFilter ();
312+ }
313+
290314static PyObject * PyFilter (PyObject * self , PyObject * args )
291315{
292316 PyObject * obj = NULL ;
293317 const char * query = NULL ;
294318
295- if (!PyArg_ParseTuple (args , "Os" , & obj , & query )) {
319+ if (!initFilter (args , & obj , & query ))
296320 return NULL ;
297- }
298-
299- itemAtAllocsList = PyList_New (0 );
300- stringAllocsList = PyList_New (0 );
301321
302322 Nodes * nodes = UastFilter (ctx , obj , query );
303323 if (!nodes ) {
304- char * error = LastError ();
305- PyErr_SetString (PyExc_RuntimeError , error );
306- free (error );
307- Py_DECREF (stringAllocsList );
308- Py_DECREF (itemAtAllocsList );
324+ filterError ();
309325 return NULL ;
310326 }
311327 size_t len = NodesSize (nodes );
@@ -320,14 +336,71 @@ static PyObject *PyFilter(PyObject *self, PyObject *args)
320336 PyObject * iter = PySeqIter_New (list );
321337 Py_DECREF (list );
322338
323- Py_DECREF (itemAtAllocsList );
324- Py_DECREF (stringAllocsList );
339+ cleanupFilter ();
325340 return iter ;
326341}
327342
343+ static PyObject * PyFilterBool (PyObject * self , PyObject * args )
344+ {
345+ PyObject * obj = NULL ;
346+ const char * query = NULL ;
347+
348+ if (!initFilter (args , & obj , & query ))
349+ return NULL ;
350+
351+ bool ok ;
352+ bool res = UastFilterBool (ctx , obj , query , & ok );
353+ if (!ok ) {
354+ filterError ();
355+ return NULL ;
356+ }
357+
358+ cleanupFilter ();
359+ return res ? Py_True : Py_False ;
360+ }
361+
362+ static PyObject * PyFilterNumber (PyObject * self , PyObject * args )
363+ {
364+ PyObject * obj = NULL ;
365+ const char * query = NULL ;
366+
367+ if (!initFilter (args , & obj , & query ))
368+ return NULL ;
369+
370+ bool ok ;
371+ double res = UastFilterNumber (ctx , obj , query , & ok );
372+ if (!ok ) {
373+ filterError ();
374+ return NULL ;
375+ }
376+
377+ cleanupFilter ();
378+ return PyFloat_FromDouble (res );
379+ }
380+
381+ static PyObject * PyFilterString (PyObject * self , PyObject * args )
382+ {
383+ PyObject * obj = NULL ;
384+ const char * query = NULL ;
385+
386+ if (!initFilter (args , & obj , & query ))
387+ return NULL ;
388+
389+ const char * res = UastFilterString (ctx , obj , query );
390+ if (res == NULL ) {
391+ filterError ();
392+ return NULL ;
393+ }
394+
395+ cleanupFilter ();
396+ return PyUnicode_FromString (res );
397+ }
328398static PyMethodDef extension_methods [] = {
329399 {"filter" , PyFilter , METH_VARARGS , "Filter nodes in the UAST using the given query" },
330400 {"iterator" , PyUastIter_new , METH_VARARGS , "Get an iterator over a node" },
401+ {"filter_bool" , PyFilterBool , METH_VARARGS , "For queries returning boolean values" },
402+ {"filter_number" , PyFilterNumber , METH_VARARGS , "For queries returning boolean values" },
403+ {"filter_string" , PyFilterString , METH_VARARGS , "For queries returning boolean values" },
331404 {NULL , NULL , 0 , NULL }
332405};
333406
0 commit comments