Tree-sitter 以及它的 Query 功能

MartinLwx 收录于类别 Program-Analysis

2024-09-07 2024-09-07 约 2987 字预计阅读 6 分钟

引言

Tree-sitter 是一个 Parse Generator，也就是用来生成 Parser 的。除此之外，它还提供了一些额外的功能，比如今天要聊到的 Tree-sitter Query ，Query 提供了一套基于 S 表达式的 DSL（Domain Specific Language），可以查询 AST，获得你想要的信息，在正式学习如何使用 Query 前，我们先讲一些相关的背景知识，好让这个文章尽量是 Self-contained 的

S 表达式

信息

如果你曾经写过 Lisp 或者是其他 Lisp 方言（比如 Scheme、Racket 等），你对 S 表达式应该很熟悉

S 表达式的定义是递归式的，定义如下

S 表达式为 x，这里的 x 是 Atom，Atom 就是不可再分解的意思
S 表达式不为空，而是 (S1 S2 ... Sn) 的形式，其中从 S1 到 Sn 都是 S 表达式

S 表达式是递归式定义的，树也是递归式定义的，因此 S 表达式可以用来表示树，以一个简单的 AST 为例

python

上面的 AST 就是一个简单的 1 + (2 - 3)，而它的 S 表达式是

lisp

(+ 1 (- 2 3))

这种情况下，每个 () 里面，第一个是父节点，其他的则是这个父节点的孩子节点

AST in Tree-sitter

Tree-sitter 支持生成很多编程语言的 Parser，每一个支持的编程语言都有一个对应的 tree-sitter-* GitHub 仓库，比如 Python 的就是 tree-sitter-python

每一个 tree-sitter-* 仓库里面都会有一个 grammar.js 文件，这个文件定义了对应编程语言的 AST 结构，以 tree-sitter-python 为例，Python 里面的 for 循环语句的语法是这样子的

for_statement: $ => seq(
  optional('async'),
  'for',
  field('left', $._left_hand_side),
  'in',
  field('right', $._expressions),
  ':',
  field('body', $._suite),
  field('alternative', optional($.else_clause)),
),

如果对编译器有一定了解，我们不难就不难看出 Python 的 For 循环语句大概是这样子的

python

[async] for <left> in <right>:
    <body>
[<else_clause>]

技巧

[] 表示可选项，<> 表示还有另外的 Non-terminal，可以进一步展开

知道 Tree-sitter 的 AST 是如何定义出来之后，我们当然要知道如何用 Tree-sitter 生成的 Parser 去生成代码的 AST

我们可以使用命令行，首先用 homebrew 安装并进行配置

$ brew install tree-sitter
$ tree-sitter init-config

然后，克隆你要分析的编程语言的 tree-sitter-* 仓库并生成 Parser，后面要以 Python 代码为例子，因此这里就用 tree-sitter-python 作为例子

$ git clone git@github.com:tree-sitter/tree-sitter-python.git
$ tree-sitter generate  # generate a parser in src/parser.c

技巧

如果想要在任意目录都可以用 tree-sitter 解析，那么就

调用 tree-sitter init-config 确定你使用的操作系统的 tree-sitter 配置文件位置
打开配置文件查看 parse-directories 字段看都有哪一些文件夹
将 tree-sitter-* 仓库放到对应文件夹下（任意挑一个）即可

现在在同个目录下，你就可以生成 Python 代码的 AST 了，让我们先写一段简单的 Python 代码

python

def sum(x: int, y: int) -> int:
    return x + y


sum(1, 2)

通过如下的命令就可以得到 S 表达式的 AST 了

$ tree-sitter parse foo.py

lisp

(module [0, 0] - [5, 0]
  (function_definition [0, 0] - [1, 16]
    name: (identifier [0, 4] - [0, 7])
    parameters: (parameters [0, 7] - [0, 23]
      (typed_parameter [0, 8] - [0, 14]
        (identifier [0, 8] - [0, 9])
        type: (type [0, 11] - [0, 14]
          (identifier [0, 11] - [0, 14])))
      (typed_parameter [0, 16] - [0, 22]
        (identifier [0, 16] - [0, 17])
        type: (type [0, 19] - [0, 22]
          (identifier [0, 19] - [0, 22]))))
    return_type: (type [0, 27] - [0, 30]
      (identifier [0, 27] - [0, 30]))
    body: (block [1, 4] - [1, 16]
      (return_statement [1, 4] - [1, 16]
        (binary_operator [1, 11] - [1, 16]
          left: (identifier [1, 11] - [1, 12])
          right: (identifier [1, 15] - [1, 16])))))
  (expression_statement [4, 0] - [4, 9]
    (call [4, 0] - [4, 9]
      function: (identifier [4, 0] - [4, 3])
      arguments: (argument_list [4, 3] - [4, 9]
        (integer [4, 4] - [4, 5])
        (integer [4, 7] - [4, 8])))))

其中 [x1, y1] - [x2, y2] 这个的意思是，这个 AST 节点对应的源码在源文件的 x1 行到 x2 行，y1 到 y2 列，索引从 0 开始，左闭右开

技巧

Tree-sitter 官方提供了一个 Playground，可以直接在线尝试，最好用的功能是你鼠标点击某个 AST 节点后，可以高亮对应的源代码，很直观

Basic Query Syntax

那么如何在 AST 中找到自己感兴趣的节点呢？这就是 Query 要做的事情了

Query 就是一个 S 表达式，比如我们想要找到 Python 代码里面的函数定义，它对应的 S 表达式如下

lisp

(function_definition)

同时，Tree-sitter 还允许我们 返回被匹配到的内容，保存在一个变量里面，只需要在 @variable_name 放在 AST 节点对应的 S 表达式的后面，比如

lisp

(function_definition) @func

技巧

现在，被匹配到的 Python 函数的源代码都会被保存在 func 这个变量里面，你可以在 Tree-sitter 官方的 Playground 里面看到匹配的结果

技巧

在 Tree-sitter 的语法文件里面，AST 节点分为 2 种，一种是带名字的（Named AST Nodes），一种是匿名的（Anonymous AST Nodes）。对于 Named AST Node，它在 AST 里面的形式是：field_name: S-expression

另外，Tree-sitter 还支持指定 AST 节点必须有某些类型的孩子节点，只需要加上 field: 这种前缀。比如我们要找有返回类型声明的 Python 函数

lisp

(function_definition
  return_type: (type)) @has_return_type_func

也支持指定 AST 节点不能有某些类型的孩子节点，只需要在 field 的前面用 ! 即可。比如下面的例子要找的是没有返回类型声明的 Python 函数

lisp

(function_definition
  !return_type) @no_return_type_func

那么对于 Anonymous AST 节点呢？直接指定要匹配的结果即可，不需要用 ()。还是用前面的代码作为例子，我们想要找二元操作，而且执行的是 + 法，就可以这么写

lisp

(binary_operator
  operator: "+") @add_op

技巧

匿名 AST 节点不会在导出的 AST 里面，那要如何写 Query 呢？查看对应编程语言的 grammar.js 文件，比如我这里就参考了这几行

Advanced Query Syntax

Query 还支持更多的语法，像 [[Regex Expressions|正则表达式]] 里面的 *, ?, +, [abc] 都支持的，考虑这么一个场景，我们有如下的 3 个 Python 列表

python

arr1 = []
arr2 = [1]
arr3 = [1, 2]

我们想要提取每个这样的列表里面的数字，但可能有 0 个，也可能有任意个，这恰恰是正则表达式里面的 * 的用法，因此我们可以写出如下的 Query

lisp

((integer) ("," (integer))*)@int_list

另外 Query 还支持 . 这种标记，根据出现的位置不同含义也不同，前面提到，对于一个 (S1 S2 ... Sn) 这样的 S 表达式，S1 是父节点，S2 ... Sn 都是它的子节点

如果 . 出现在 S1 后面，那么 S2 必须是父节点 S1 的第一个孩子。现在假设我们要从 arr1, arr2, arr3 中提取第一个元素（若有），那么我们可以这么写

lisp

(list . (integer) @first_int)

如果 . 出现在 Sn 的后面，那么 Sn 必须是父节点 S1 的最后一个孩子，同理，我们可以提取 arr1, arr2, arr3 的最后一个元素

lisp

(list (integer) @last_int .)

现在情况更加负载，Python 列表是动态的，可能包含各种数据类型的元素，比如

python

arr1 = []
arr2 = [None]
arr3 = [1, 2, "90"]

这个时候我们就不能用 (integer) 去匹配了，我们需要一种特殊的表达式，表示可以匹配任意 AST 节点类型，在 Python 里面，_ 可以起到这样的功能。Tree-sitter 的 Query 也是，所以我们可以写

lisp

(list (_) @last_int .)

技巧

(_) 可以匹配 Named AST Node，_ 可以匹配 Named AST Node 或者 Anonymous AST Node

More Advanced Query Syntax

只是用 @variable_name 获取被匹配的内容似乎没有多大用处，要是可以基于被匹配的内容做进一步的筛选就好了，Tree-sitter Query 也是支持的，这叫做 predicate

每一个 predicate 也是一个 S 表达式，形如 (S1 S2 S3)

S1 的位置可以是 #eq?, #not-eq?, #any-eq?, #any-not-eq?, #match?, #not-match?, #any-match?, #any-not-match?, #any-of?, #not-any-of? 中的任意一个，看名字就知道意思了
S2 的位置必须是 @variable_name
S3 的位置可以是 @variable_name，或者是字符串（使用 *eq? 时）、正则表达式（使用 *match? 时）、候选列表（使用 *-of? 时，* 直接列出来即可，不需要用 () 括起来）

信息

Lisp 或者是其他的 Lisp 方言里面就有一个 eq? 函数

现在，我们尝试用 predicate 进行筛选：提取列表里面的最后一个元素，而且当且仅当它是字符串 "90" 的时候，不难用下面的 Query 表达式进行验证

lisp

(list (string) @last_str (#eq? @last_str "90") .)

Bindings

前面都是用命令行工具或者直接在官网的 Playground 上进行验证，那能否将这些功能融入其他编程语言的代码里面呢？当然可以，具体处置的编程语言可以看官网

以 Python 的 binding 为例，我们需要先安装

$ pip install tree-sitter
$ pip install tree-sitter-python

然后就可以开始编写代码了，先生成 Parser 然后解析代码得到 AST

python

import tree_sitter_python as tspython
from tree_sitter import Language, Parser

PY_LANGUAGE = Language(tspython.language())

code = """def sum(x: int, y: int) -> int:
    return x + y


sum(1, 2)
"""

parser = Parser(PY_LANGUAGE)
tree = parser.parse(bytes(code, encoding="utf8"))

现在 tree 就是解析得到的 AST，我们可以查看根节点的类型以及对应的源码的起始位置

python

parser = Parser(PY_LANGUAGE)
tree = parser.parse(bytes(code, encoding="utf8"))
root_node = tree.root_node

print(root_node.type)  # module
print(root_node.start_point)  # Point(row=0, column=0)
print(root_node.end_point)  # Point(row=5, column=0)

可以看到，结果和之前命令行的输出是一致的

lisp

(module [0, 0] - [5, 0]
... ;; omit

这里我们可以写一个 Query，提取函数里面带类型提示的入参，并打印参数名、参数类型

python

query = PY_LANGUAGE.query(
    """(typed_parameter
         (identifier)@name
         type: (type (identifier)@val))"""
)
captures = query.captures(root_node)
for arg_name, arg_type in zip(captures["name"], captures["val"]):
    print(f"{str(arg_name.text)}: {str(arg_type.text)}")

总结

Tree-sitter 是一个强大又轻量的 Parser Generator，在做程序分析的时候我经常都会用它生成的 Parser 加上 Query 功能在 AST 里面提取信息，这些信息可以用于做进一步的分析。和 ANTLR 以及其他的 Parser Generator 相比，我感觉还是 Tree-sitter 更好用一些。这可能是因为我学过 Scheme、Racket 这种 Lisp 形式的编程语言，用 S 表达式写 Query 让我感觉很亲切 :)