esquery + python

esquery.py

import re

from esprima.visitor import NodeVisitor, Visited

rsel = re.compile("(?:(?P<type>\\w+)|\\s+|(?P<child>>))")


def parse_selector(sel):
    matches = [m.groupdict() for m in rsel.finditer(sel)]
    parsed = []
    for m in matches:
        if m["type"]:
            parsed.append({"kind": "NodeType", "type": m["type"]})
        elif m["child"]:
            parsed.append({"kind": "Child"})
    return parsed


def esquery(node, selector):
    parsed = parse_selector(selector)
    target = parsed[-1]
    if target["kind"] != "NodeType":
        raise Exception("Selector syntax error TODO")

    ancestry = []
    target_matches = []

    def visit_Object(obj):
        ancestry.append(obj)
        if obj.type == target["type"]:
            target_matches.append(ancestry[:])
        yield obj.__dict__
        yield Visited(obj)
        ancestry.pop()

    visitor = NodeVisitor()
    visitor.visit_Object = visit_Object

    visitor.visit(node)

    matches = []
    for tm in target_matches:
        ptr = 0
        node = None

        while ptr < len(parsed):
            selector = parsed[ptr]
            ptr += 1

            if selector["kind"] == "NodeType":
                try:
                    node = next(node for node in tm if node.type == selector["type"])
                    tm = tm[tm.index(node) + 1 :]
                except StopIteration:
                    node = None
                    break

            elif selector["kind"] == "Child":
                next_selector = parsed[ptr]
                ptr += 1
                node = tm[0]
                if node.type == next_selector["type"]:
                    tm = tm[1:]
                else:
                    node = None
                    break

        if node is not None:
            matches.append(node)

    return matches