python环境解析任意编程语言 tree-sitter使用方法（1）

背景

我个人目前仍在研究代码有关的知识。目前基于深度学习表征代码的论文越来越卷了，用到的工具越来越高级了。目前有一个开源项目tree-sitter，专门用于解析具体语法树，声称：

1. 足够通用，能用于任何编程语言
2. 足够迅速，能在文本编辑器中响应每一个用户输入
3. 足够鲁棒，即便语法错误也能解析语法树
4. 无依赖性，能很好地嵌入于程序中

在官方提供的playground玩了玩，的确1、2、3点都很符合。
所以个人做 （水）了本篇文章。

安装

py-tree-sitter已经做了详细的描述，所以这里简短描述，顺便说个遇到的问题。

1. 找个合适的python环境，install

pip3 install tree_sitter

2. 对于要解析的编程语言，随便创建文件夹（比如vendor），该目录git clone指定语言的仓库，在tree-sitter官网这里找，比如我对Java、Python、C++、C#、JS感兴趣：

git clone https://github.com/tree-sitter/tree-sitter-java
git clone https://github.com/tree-sitter/tree-sitter-python
git clone https://github.com/tree-sitter/tree-sitter-cpp
git clone https://github.com/tree-sitter/tree-sitter-c-sharp
git clone https://github.com/tree-sitter/tree-sitter-javascript

需要注意的是，C++对应cpp，C#对应c-sharp，后面使用的时候需要认清楚官方定义的名称。

3. 创建build文件夹，用于保存xxx.so文件，该文件相当于自定义的编译器，用于解析代码生成语法树。然后复制以下代码运行。

from tree_sitter import Language

Language.build_library(
  # so文件保存位置
  'build/my-languages.so',

  # vendor文件下git clone的仓库
  [
    'vendor/tree-sitter-java',
    'vendor/tree-sitter-python',
    'vendor/tree-sitter-cpp',
    'vendor/tree-sitter-c-sharp',
    'vendor/tree-sitter-javascript',
  ]
)

这里有一个小插曲，个人用windows电脑，一开始运行这段代码直接报错，好像说缺少什么msvc文件，所以我还下载了visual studio才解决。现在看到tree-sitter的__init__.py文件下，有一条compiler = new_compiler()代码，发现以下代码：

if compiler is None:
	# get_default_compiler 用于选择_default_compilers：
	'''
	_default_compilers = (
	    ('cygwin.*', 'unix'),
	    ('posix', 'unix'),
	    ('nt', 'msvc'),
    )
    '''
    compiler = get_default_compiler(plat) # windows是nt

'''
compiler_class = { 'unix':    ('unixccompiler', 'UnixCCompiler',
                               "standard UNIX-style compiler"),
                   'msvc':    ('_msvccompiler', 'MSVCCompiler',
                               "Microsoft Visual C++"),
                   'cygwin':  ('cygwinccompiler', 'CygwinCCompiler',
                               "Cygwin port of GNU C Compiler for Win32"),
                   'mingw32': ('cygwinccompiler', 'Mingw32CCompiler',
                               "Mingw32 port of GNU C Compiler for Win32"),
                   'bcpp':    ('bcppcompiler', 'BCPPCompiler',
                               "Borland C++ Compiler"),
                 }
'''
(module_name, class_name, long_description) = compiler_class[compiler]

看了代码后就清楚了，我之前电脑缺少Microsoft Visual C++，安装visual studio，配置C++后就好了。

4. 解析

from tree_sitter import Language, Parser

# 注意C++对应cpp，C#对应c-sharp
# 看仓库名称
CPP_LANGUAGE = Language('build/my-languages.so', 'cpp')
CS_LANGUAGE = Language('build/my-languages.so', 'c-sharp')

# 举一个CPP例子
cpp_parser = Parser()
cpp_parser.set_language(CPP_LANGUAGE)

# 这是b站网友写的代码，解析看看
cpp_code_snippet = '''
int mian{
  piantf("hello word");
  remake O;
}
'''

# 没报错就是成功
tree = cpp_parser.parse(bytes(cpp_code_snippet, "utf8"))
# 注意，root_node 才是可遍历的树节点
root_node = tree.root_node

解析小例子

我发现这个tree-sitter库是看到文章GraphCodeBert了解到，后来，很多研究比如UniXcoder，CodeT5，TreeBert和SynCoBert【不开源的文章】等等都用了该库。
【吐槽：深度学习表征代码越来越卷了，Money and Equipment Is All You Need 属于是了，各个下游任务刷榜。本来实验室刚从传统算法转机器学习，就一个GPU，留给我硕士菜鸡的毕业的机会都快弄没了。】

GraphCodeBert使用语法树分词的方法还是不错的，这里是原论文别人写的代码，网址在这里，个人觉得很不错，供参考：

from tree_sitter import Language, Parser


def tree_to_token_index(root_node):
    if (len(root_node.children) == 0 or root_node.type == 'string') and root_node.type != 'comment':
        return [(root_node.start_point, root_node.end_point)]
    else:
        code_tokens = []
        for child in root_node.children:
            code_tokens += tree_to_token_index(child)
        return code_tokens


def index_to_code_token(index, code):
    start_point = index[0]
    end_point = index[1]
    if start_point[0] == end_point[0]:
        s = code[start_point[0]][start_point[1]:end_point[1]]
    else:
        s = ""
        s += code[start_point[0]][start_point[1]:]
        for i in range(start_point[0]+1, end_point[0]):
            s += code[i]
        s += code[end_point[0]][:end_point[1]]
    return s


if __name__ == '__main__':
    CPP_LANGUAGE = Language('build/my-languages.so', 'cpp')
    cpp_parser = Parser()
    cpp_parser.set_language(CPP_LANGUAGE)

    cpp_code_snippet = '''
    int mian{
    piantf("hello word");
    remake O;
    }
    '''

    tree = cpp_parser.parse(bytes(cpp_code_snippet, "utf8"))
    root_node = tree.root_node
    tokens_index = tree_to_token_index(root_node)
    cpp_loc = cpp_code_snippet.split('\n')
    code_tokens = [index_to_code_token(x, cpp_loc) for x in tokens_index]
    print(code_tokens)

提取具体节点

其实tree-sitter还可以手动配置想要的语法树节点，通过定义query，便于直接提取特定语法树节点。我看有代码为了定位语法树节点，dfs语法树，手动写判断，一大堆代码，还要回溯判断父节点，太困难了。等一段时间后，在接下来的文章中再介绍。