juanjux · Feb 21, 2018 · Feb 21, 2018
@@ -52,6 +52,12 @@ print(uast)
 # "filter' allows you to use XPath queries to filter on result nodes:
 print(bblfsh.filter(uast, "//Import[@roleImport and @roleDeclaration]//alias")
 
+# filter\_[bool|string|number] must be used when using XPath functions returning
+# these types:
+print(bblfsh.filter_bool(uast, "boolean(//*[@strtOffset or @endOffset])"))
+print(bblfsh.filter_string(uast, "name(//*[1])"))
+print(bblfsh.filter_number(uast, "count(//*)"))
+
 # You can also iterate on several tree iteration orders:
 it = bblfsh.iterator(uast, bblfsh.TreeOrder.PRE_ORDER)
 for node in it:

@@ -1,5 +1,5 @@
 from bblfsh.client import BblfshClient
-from bblfsh.pyuast import filter, iterator
+from bblfsh.pyuast import filter, filter_bool, filter_number, filter_string, iterator
 from bblfsh.aliases import *
 
 class TreeOrder:

@@ -39,7 +39,6 @@ static size_t Size(const void *node, const char *prop) {
   PyObject *o = Attribute(node, prop);
   if (o != NULL) {
     retval = PySequence_Size(o);
-    Py_DECREF(o);
   }
 
   return retval;
@@ -287,25 +286,42 @@ static void PyUastIter_dealloc(PyObject *self)
   UastIteratorFree(((PyUastIter *)self)->iter);
 }
 
+static bool initFilter(PyObject *args, PyObject **obj, const char **query)
+{
+  if (!PyArg_ParseTuple(args, "Os", obj, query)) {
+    return false;
+  }
+
+  itemAtAllocsList = PyList_New(0);
+  stringAllocsList = PyList_New(0);
+  return true;
+}
+
+static void cleanupFilter(void)
+{
+    Py_DECREF(itemAtAllocsList);
+    Py_DECREF(stringAllocsList);
+}
+
+static void filterError(void)
+{
+  char *error = LastError();
+  PyErr_SetString(PyExc_RuntimeError, error);
+  free(error);
+  cleanupFilter();
+}
+
 static PyObject *PyFilter(PyObject *self, PyObject *args)
 {
   PyObject *obj = NULL;
   const char *query = NULL;
 
-  if (!PyArg_ParseTuple(args, "Os", &obj, &query)) {
+  if (!initFilter(args, &obj, &query))
     return NULL;
-  }
-
-  itemAtAllocsList = PyList_New(0);
-  stringAllocsList = PyList_New(0);
 
   Nodes *nodes = UastFilter(ctx, obj, query);
   if (!nodes) {
-    char *error = LastError();
-    PyErr_SetString(PyExc_RuntimeError, error);
-    free(error);
-    Py_DECREF(stringAllocsList);
-    Py_DECREF(itemAtAllocsList);
+    filterError();
     return NULL;
   }
   size_t len = NodesSize(nodes);
@@ -320,14 +336,71 @@ static PyObject *PyFilter(PyObject *self, PyObject *args)
   PyObject *iter = PySeqIter_New(list);
   Py_DECREF(list);
 
-  Py_DECREF(itemAtAllocsList);
-  Py_DECREF(stringAllocsList);
+  cleanupFilter();
   return iter;
 }
 
+static PyObject *PyFilterBool(PyObject *self, PyObject *args)
+{
+  PyObject *obj = NULL;
+  const char *query = NULL;
+
+  if (!initFilter(args, &obj, &query))
+    return NULL;
+
+  bool ok;
+  bool res = UastFilterBool(ctx, obj, query, &ok);
+  if (!ok) {
+    filterError();
+    return NULL;
+  }
+
+  cleanupFilter();
+  return res ? Py_True : Py_False;
+}
+
+static PyObject *PyFilterNumber(PyObject *self, PyObject *args)
+{
+  PyObject *obj = NULL;
+  const char *query = NULL;
+
+  if (!initFilter(args, &obj, &query))
+    return NULL;
+
+  bool ok;
+  double res = UastFilterNumber(ctx, obj, query, &ok);
+  if (!ok) {
+    filterError();
+    return NULL;
+  }
+
+  cleanupFilter();
+  return PyFloat_FromDouble(res);
+}
+
+static PyObject *PyFilterString(PyObject *self, PyObject *args)
+{
+  PyObject *obj = NULL;
+  const char *query = NULL;
+
+  if (!initFilter(args, &obj, &query))
+    return NULL;
+
+  const char *res = UastFilterString(ctx, obj, query);
+  if (res == NULL) {
+    filterError();
+    return NULL;
+  }
+
+  cleanupFilter();
+  return PyUnicode_FromString(res);
+}
 static PyMethodDef extension_methods[] = {
     {"filter", PyFilter, METH_VARARGS, "Filter nodes in the UAST using the given query"},
     {"iterator", PyUastIter_new, METH_VARARGS, "Get an iterator over a node"},
+    {"filter_bool", PyFilterBool, METH_VARARGS, "For queries returning boolean values"},
+    {"filter_number", PyFilterNumber, METH_VARARGS, "For queries returning boolean values"},
+    {"filter_string", PyFilterString, METH_VARARGS, "For queries returning boolean values"},
     {NULL, NULL, 0, NULL}
 };
 

@@ -4,7 +4,8 @@
 import docker
 
 from bblfsh import (BblfshClient, filter, iterator, role_id,
-        role_name, Node, ParseResponse, TreeOrder)
+        role_name, Node, ParseResponse, TreeOrder, filter_bool,
+        filter_number, filter_string)
 from bblfsh.launcher import ensure_bblfsh_is_running
 from bblfsh.client import NonUTF8ContentException
 
@@ -28,10 +29,10 @@ def setUp(self):
 
     def testVersion(self):
         version = self.client.version()
-        assert(hasattr(version, "version"))
-        assert(version.version)
-        assert(hasattr(version, "build"))
-        assert(version.build)
+        self.assertTrue(hasattr(version, "version"))
+        self.assertTrue(version.version)
+        self.assertTrue(hasattr(version, "build"))
+        self.assertTrue(version.build)
 
     def testNativeParse(self):
         reply = self.client.native_parse(__file__)
@@ -119,13 +120,33 @@ def testFilterEndCol(self):
         self.assertTrue(any(filter(node, "//*[@endCol=50]")))
         self.assertFalse(any(filter(node, "//*[@endCol=5]")))
 
+    def testFilterBool(self):
+        node = Node()
+        self.assertTrue(filter_bool(node, "boolean(//*[@startOffset or @endOffset])"))
+        self.assertFalse(filter_bool(node, "boolean(//*[@blah])"))
+
+    def testFilterNumber(self):
+        node = Node()
+        node.children.extend([Node(), Node(), Node()])
+        self.assertEqual(int(filter_number(node, "count(//*)")), 4)
+
+    def testFilterString(self):
+        node = Node()
+        node.internal_type = "test"
+        self.assertEqual(filter_string(node, "name(//*[1])"), "test")
+
     def testFilterBadQuery(self):
         node = Node()
         self.assertRaises(RuntimeError, filter, node, "//*roleModule")
 
+    def testFilterBadType(self):
+        node = Node()
+        node.end_position.col = 50
+        self.assertRaises(RuntimeError, filter, node, "boolean(//*[@startPosition or @endPosition])")
+
     def testRoleIdName(self):
-        assert(role_id(role_name(1)) == 1)
-        assert(role_name(role_id("IDENTIFIER")) == "IDENTIFIER")
+        self.assertEqual(role_id(role_name(1)), 1)
+        self.assertEqual(role_name(role_id("IDENTIFIER")),  "IDENTIFIER")
 
     def _itTestTree(self):
         root = Node()
@@ -186,7 +207,7 @@ def _validate_resp(self, resp):
         # self.assertIsInstance(resp.uast, Node)
 
         # Sometimes its fully qualified, sometimes is just "Node"... ditto
-        assert(resp.uast.__class__.__name__.endswith('Node'))
+        self.assertTrue(resp.uast.__class__.__name__.endswith('Node'))
 
     def testManyFilters(self):
         root = self.client.parse(__file__).uast

@@ -134,7 +134,7 @@ def main():
         },
         name="bblfsh",
         description="Fetches Universal Abstract Syntax Trees from Babelfish.",
-        version="2.8.2",
+        version="2.9.0",
         license="Apache 2.0",
         author="source{d}",
         author_email="[email protected]",