Tree-sitter & Jedi

import tree_sitter_python as tspython from tree_sitter import Language, Parser parser = Parser(Language(tspython.language())) code = """ def greet(name: str) -> str: return f"Hello, {name}!" class Bot: def __init__(self, model): self.model = model """ tree = parser.parse(bytes(code, "utf-8")) root = tree.root_node # 根节点，类型是 module

def extract_definitions(node): """递归遍历 AST，提取所有函数和类定义""" results = [] for child in node.children: if child.type == "function_definition": name = child.child_by_field_name("name").text.decode() params = child.child_by_field_name("parameters") ret = child.child_by_field_name("return_type") results.append({ "type": "function", "name": name, "params": params.text.decode() if params else "", "return_type": ret.text.decode() if ret else None, "start_line": child.start_point[0], "end_line": child.end_point[0], "body_text": child.text.decode(), # 完整函数代码 }) elif child.type == "class_definition": name = child.child_by_field_name("name").text.decode() supers = child.child_by_field_name("superclasses") body = child.child_by_field_name("body") # 递归提取类内部的方法 methods = extract_definitions(body) if body else [] results.append({ "type": "class", "name": name, "bases": supers.text.decode() if supers else "", "methods": methods, }) # 递归进入子节点（处理嵌套定义） results.extend(extract_definitions(child)) return results

# Query 基本语法：S-表达式 # 匹配所有函数定义，并捕获函数名 query_str = """ (function_definition name: (identifier) @func.name parameters: (parameters) @func.params return_type: (type)? @func.return_type body: (block) @func.body ) @func.def """ # 使用 Query query = Language(tspython.language()).query(query_str) captures = query.captures(tree.root_node) # captures 是 dict: {"func.name": [node, ...], "func.def": [node, ...], ...} for node in captures.get("func.name", []): print(node.text.decode()) # 输出每个函数名

import jedi # 创建 Script —— Jedi 的核心入口 script = jedi.Script( source="import os\nos.path.join('a', 'b')", path="example.py", # 可选：文件路径（帮助解析相对导入） ) # 或者直接从文件创建 script = jedi.Script(path="my_module.py")

import jedi code = """ from pathlib import Path def read_config(filepath: str) -> dict: p = Path(filepath) return json.loads(p.read_text()) result = read_config("config.json") """ script = jedi.Script(source=code, path="app.py") # goto: "read_config" 在第 7 行被调用，跳转到定义处 defs = script.goto(7, 10) # → [Name: read_config, line=3] # infer: "result" 的类型是什么？ types = script.infer(7, 1) # → [Name: dict] # references: "Path" 在哪些地方被使用？ refs = script.references(2, 21) # → [line 2 (import), line 5 (usage)] # get_names: 获取文件所有顶层名称 names = script.get_names() for n in names: print(n.name, n.type, n.line) # Path module 2 # read_config function 3 # result statement 7

import jedi def build_function_info(filepath): """用 Jedi 提取文件中所有函数的语义信息""" script = jedi.Script(path=filepath) names = script.get_names(all_scopes=False) # 仅顶层 functions = [] for name in names: if name.type != "function": continue func_info = { "name": name.name, # "read_config" "full_name": name.full_name, # "app.read_config" "description": name.description, # "def read_config(filepath: str) -> dict" "module": name.module_name, # "app" "line": name.line, # 3 "is_def": name.is_definition(),# True } # 获取函数的参数（子名称） params = [n.name for n in name.defined_names() if n.type == "param"] func_info["params"] = params # 用 goto 追踪函数内部的调用 # (通常配合 Tree-sitter 找到调用位置后再用 Jedi) functions.append(func_info) return functions

import jedi import tree_sitter_python as tspython from tree_sitter import Language, Parser from pathlib import Path PY_LANGUAGE = Language(tspython.language()) parser = Parser(PY_LANGUAGE) # ── Step 1: Tree-sitter 提取结构 ── def extract_structure(filepath): """用 Tree-sitter 快速提取文件的代码结构""" code = Path(filepath).read_bytes() tree = parser.parse(code) # Query: 一次性提取函数、类、导入、调用 q = PY_LANGUAGE.query(""" (function_definition name: (identifier) @func_name) @func (class_definition name: (identifier) @class_name) @class (import_from_statement module_name: (dotted_name) @import_source name: (dotted_name) @import_name) @import (call function: (identifier) @call_name) @call_expr (call function: (attribute attribute: (identifier) @method_name)) @method_call """) captures = q.captures(tree.root_node) return captures # ── Step 2: Jedi 补充语义 ── def enrich_with_semantics(filepath, ts_captures): """用 Jedi 为 Tree-sitter 提取的节点添加语义信息""" script = jedi.Script(path=filepath) graph_nodes = [] graph_edges = [] # 处理函数定义 for node in ts_captures.get("func_name", []): # Tree-sitter 行号从 0 开始，Jedi 从 1 开始 line = node.start_point[0] + 1 col = node.start_point[1] # 用 Jedi 获取完全限定名和类型 jedi_names = script.goto(line, col) if jedi_names: n = jedi_names[0] graph_nodes.append({ "id": n.full_name or n.name, "type": "function", "name": n.name, "desc": n.description, "file": str(filepath), "line": n.line, }) # 处理调用 —— 构建"调用"边 for node in ts_captures.get("call_name", []): line = node.start_point[0] + 1 col = node.start_point[1] # Jedi 追踪到定义处 targets = script.goto(line, col) if targets: t = targets[0] # 找到调用者（包含此调用的函数） caller = find_enclosing_function(node) # 用 Tree-sitter parent 向上找 if caller: graph_edges.append({ "from": caller, "to": t.full_name or t.name, "relation": "calls", }) return graph_nodes, graph_edges def find_enclosing_function(node): """沿 Tree-sitter AST 向上查找包含此节点的函数""" current = node.parent while current: if current.type == "function_definition": name_node = current.child_by_field_name("name") return name_node.text.decode() if name_node else None current = current.parent return None