Cyclic0007/ruff
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Write an extractor to generate AST from tree-sitter

Browse Source
This commit is contained in:
Charlie Marsh
2022-10-05 17:51:45 -04:00
parent 40ab5d353b
commit c8dad90115
3 changed files with 561 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

562
examples/parse_python.rs
View File

@@ -1,10 +1,431 @@
use anyhow::Result;
use tree_sitter::{Parser, Query, QueryCursor};
extern crate core;
enum Action {
Up,
Down,
Right,
use anyhow::Result;
use num_bigint::BigInt;
use num_traits::{float, Num};
use rustpython_ast::{
Arguments, Constant, Expr, ExprContext, ExprKind, Keyword, KeywordData, Location, Operator,
Stmt, StmtKind, Withitem,
};
use tree_sitter::{Node, Parser, Point};
fn to_location(point: Point) -> Location {
Location::new(point.row + 1, point.column + 1)
}
fn print_node(node: Node, source: &[u8]) {
let range = node.range();
let text = &source[range.start_byte..range.end_byte];
let line = range.start_point.row;
let col = range.start_point.column;
println!(
"[Line: {}, Col: {}] {}: `{}`",
line,
col,
node.kind(),
std::str::from_utf8(text).unwrap()
);
}
fn extract_module(node: Node, source: &[u8]) -> Vec<Stmt> {
let mut cursor = node.walk();
node.children(&mut cursor)
.map(|node| extract_statement(node, source))
.collect()
}
fn extract_suite(node: Node, source: &[u8]) -> Vec<Stmt> {
let mut cursor = node.walk();
node.children(&mut cursor)
.map(|node| extract_statement(node, source))
.collect()
}
fn extract_text(node: Node, source: &[u8]) -> String {
let range = node.range();
let text = &source[range.start_byte..range.end_byte];
std::str::from_utf8(text).unwrap().to_string()
}
fn extract_augmented_operator(node: Node, source: &[u8]) -> Operator {
match node.kind() {
"+=" => Operator::Add,
"-=" => Operator::Sub,
"*=" => Operator::Mult,
"@=" => Operator::MatMult,
"/=" => Operator::Div,
"%=" => Operator::Mod,
"**=" => Operator::Pow,
"<<=" => Operator::LShift,
">>=" => Operator::RShift,
"|=" => Operator::BitOr,
"^=" => Operator::BitXor,
"&=" => Operator::BitAnd,
"//=" => Operator::FloorDiv,
_ => panic!("Invalid operator: {:?}", node),
}
}
fn extract_operator(node: Node, source: &[u8]) -> Operator {
match node.kind() {
"+" => Operator::Add,
"-" => Operator::Sub,
"*" => Operator::Mult,
"@" => Operator::MatMult,
"/" => Operator::Div,
"%" => Operator::Mod,
"**" => Operator::Pow,
"<<" => Operator::LShift,
">>" => Operator::RShift,
"|" => Operator::BitOr,
"^" => Operator::BitXor,
"&" => Operator::BitAnd,
"//" => Operator::FloorDiv,
_ => panic!("Invalid operator: {:?}", node),
}
}
fn extract_arguments(node: Node, source: &[u8]) -> Arguments {
Arguments {
posonlyargs: vec![],
args: vec![],
vararg: None,
kwonlyargs: vec![],
kw_defaults: vec![],
kwarg: None,
defaults: vec![],
}
}
fn extract_with_clause(node: Node, source: &[u8]) -> Vec<Withitem> {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
print_node(child, source);
}
return vec![];
}
fn extract_statement(node: Node, source: &[u8]) -> Stmt {
match node.kind() {
"for_statement" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::For {
target: Box::new(extract_expression(
node.child_by_field_name("left").unwrap(),
source,
)),
iter: Box::new(extract_expression(
node.child_by_field_name("right").unwrap(),
source,
)),
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
// STOPSHIP(charlie): Unimplemented.
orelse: vec![],
type_comment: None,
},
),
"while_statement" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::While {
test: Box::new(extract_expression(
node.child_by_field_name("condition").unwrap(),
source,
)),
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
// STOPSHIP(charlie): Unimplemented.
orelse: vec![],
},
),
"with_statement" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::With {
// TODO(charlie): If async, this will be 2? Also, we need to iterate until we find
// this, probably.
items: extract_with_clause(node.child(1).unwrap(), source),
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
type_comment: None,
},
),
"class_definition" => {
if let Some((bases, keywords)) = node
.child_by_field_name("superclasses")
.map(|node| extract_argument_list(node, source))
{
Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::ClassDef {
name: extract_text(node.child_by_field_name("name").unwrap(), source),
bases,
keywords,
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
// TODO(charlie): How do I access these? Probably need to pass them down or
// recurse.
decorator_list: vec![],
},
)
} else {
Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::ClassDef {
name: extract_text(node.child_by_field_name("name").unwrap(), source),
bases: vec![],
keywords: vec![],
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
// TODO(charlie): How do I access these? Probably need to pass them down or
// recurse.
decorator_list: vec![],
},
)
}
}
"function_definition" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::FunctionDef {
name: extract_text(node.child(1).unwrap(), source),
args: Box::new(extract_arguments(node.child(2).unwrap(), source)),
body: extract_suite(node.child_by_field_name("body").unwrap(), source),
decorator_list: vec![],
returns: None,
type_comment: None,
},
),
"return_statement" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::Return {
value: node
.child(1)
.map(|node| Box::new(extract_expression(node, source))),
},
),
"pass_statement" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::Pass,
),
"expression_statement" => {
let node = node.child(0).unwrap();
match node.kind() {
"assignment" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::Assign {
targets: vec![],
value: Box::new(extract_expression(node.child(2).unwrap(), source)),
type_comment: None,
},
),
"augmented_assignment" => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::AugAssign {
target: Box::new(extract_expression(
node.child_by_field_name("left").unwrap(),
source,
)),
value: Box::new(extract_expression(
node.child_by_field_name("right").unwrap(),
source,
)),
op: extract_augmented_operator(
node.child_by_field_name("operator").unwrap(),
source,
),
},
),
_ => Stmt::new(
to_location(node.start_position()),
to_location(node.end_position()),
StmtKind::Expr {
value: Box::new(extract_expression(node, source)),
},
),
}
}
_ => panic!("Unhandled node: {}", node.kind()),
}
}
fn extract_expression_list(node: Node, source: &[u8]) -> Vec<Expr> {
let mut cursor = node.walk();
node.children(&mut cursor)
.filter(|node| node.kind() != "(" && node.kind() != ")" && node.kind() != ",")
.map(|node| extract_expression(node, source))
.collect()
}
fn extract_keyword_argument(node: Node, source: &[u8]) -> Keyword {
Keyword::new(
Default::default(),
Default::default(),
KeywordData {
arg: Some(extract_text(
node.child_by_field_name("name").unwrap(),
source,
)),
value: Box::new(extract_expression(
node.child_by_field_name("value").unwrap(),
source,
)),
},
)
}
fn extract_argument_list(node: Node, source: &[u8]) -> (Vec<Expr>, Vec<Keyword>) {
let mut args = vec![];
let mut keywords = vec![];
for child in node.children(&mut node.walk()) {
match child.kind() {
"keyword_argument" => {
keywords.push(extract_keyword_argument(child, source));
}
"identifier" | "integer" => {
args.push(extract_expression(child, source));
}
_ => {}
}
}
(args, keywords)
}
fn extract_expression(node: Node, source: &[u8]) -> Expr {
match node.kind() {
"integer" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Int(
BigInt::from_str_radix(&extract_text(node, source), 10).unwrap(),
),
kind: None,
},
),
"float" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Float(extract_text(node, source).parse::<f64>().unwrap()),
kind: None,
},
),
"string" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Str(extract_text(node, source)),
kind: None,
},
),
"tuple" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Tuple {
elts: extract_expression_list(node, source),
ctx: ExprContext::Load,
},
),
"identifier" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Name {
id: std::str::from_utf8(&source[node.range().start_byte..node.range().end_byte])
.unwrap()
.to_string(),
ctx: ExprContext::Load,
},
),
"call" => {
let argument_list =
extract_argument_list(node.child_by_field_name("arguments").unwrap(), source);
Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Call {
func: Box::new(extract_expression(
node.child_by_field_name("function").unwrap(),
source,
)),
args: argument_list.0,
keywords: argument_list.1,
},
)
}
"binary_operator" => {
print_node(node, source);
Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::BinOp {
left: Box::new(extract_expression(
node.child_by_field_name("left").unwrap(),
source,
)),
op: extract_operator(node.child_by_field_name("operator").unwrap(), source),
right: Box::new(extract_expression(
node.child_by_field_name("right").unwrap(),
source,
)),
},
)
}
"true" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Bool(true),
kind: None,
},
),
"false" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Bool(false),
kind: None,
},
),
"ellipsis" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Constant {
value: Constant::Ellipsis,
kind: None,
},
),
"yield" => match node.child(1) {
None => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Yield { value: None },
),
Some(node) => match node.kind() {
"from" => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::YieldFrom {
value: Box::new(extract_expression(node.next_sibling().unwrap(), source)),
},
),
_ => Expr::new(
to_location(node.start_position()),
to_location(node.end_position()),
ExprKind::Yield {
value: Some(Box::new(extract_expression(node, source))),
},
),
},
},
_ => {
print_node(node, source);
panic!("Unhandled node: {}", node.kind())
}
}
}
fn main() -> Result<()> {
@@ -13,117 +434,38 @@ def double(x):
# Return a double.
return x * 2
x = double(1)
y = (f"{x}" "b")
x = (double(500), double(2, z=1))
x += 1
class Foo:
pass
for x in range(5):
yield x
yield from x
x = True
x = b"abc"
while True:
pass
with (
foo as bar,
baz as wop):
pass
"#;
let mut parser = Parser::new();
parser
.set_language(tree_sitter_python::language())
.expect("Error loading Python grammar");
let parse_tree = parser.parse(src, None);
let parse_tree = parser.parse(src.as_bytes(), None);
if let Some(parse_tree) = &parse_tree {
// Check for comments.
let query = Query::new(tree_sitter_python::language(), "(comment) @capture")?;
let mut query_cursor = QueryCursor::new();
let all_matches = query_cursor.matches(&query, parse_tree.root_node(), src.as_bytes());
for each_match in all_matches {
for capture in each_match.captures.iter() {
let range = capture.node.range();
let text = &src[range.start_byte..range.end_byte];
let line = range.start_point.row;
let col = range.start_point.column;
println!(
"[Line: {}, Col: {}] Offending source code: `{}`",
line, col, text
);
}
}
// Check for string concatenations.
let query = Query::new(
tree_sitter_python::language(),
"(concatenated_string) @capture",
)?;
let mut query_cursor = QueryCursor::new();
let all_matches = query_cursor.matches(&query, parse_tree.root_node(), src.as_bytes());
for each_match in all_matches {
for capture in each_match.captures.iter() {
let range = capture.node.range();
let text = &src[range.start_byte..range.end_byte];
let line = range.start_point.row;
let col = range.start_point.column;
println!(
"[Line: {}, Col: {}] Offending source code: `{}`",
line, col, text
);
}
}
// Walk the tree.
let mut cursor = parse_tree.walk();
let mut action = Action::Down;
loop {
match action {
Action::Up => {
if cursor.goto_next_sibling() {
action = Action::Right;
} else if cursor.goto_parent() {
action = Action::Up;
} else {
break;
}
}
Action::Down => {
let range = cursor.node().range();
let text = &src[range.start_byte..range.end_byte];
let line = range.start_point.row;
let col = range.start_point.column;
println!(
"[Line: {}, Col: {}] {}: `{}`",
line,
col,
cursor.node().kind(),
text
);
if cursor.goto_first_child() {
action = Action::Down;
} else if cursor.goto_next_sibling() {
action = Action::Right;
} else if cursor.goto_parent() {
action = Action::Up;
} else {
break;
}
}
Action::Right => {
let range = cursor.node().range();
let text = &src[range.start_byte..range.end_byte];
let line = range.start_point.row;
let col = range.start_point.column;
println!(
"[Line: {}, Col: {}] {}: `{}`",
line,
col,
cursor.node().kind(),
text
);
if cursor.goto_first_child() {
action = Action::Down;
} else if cursor.goto_next_sibling() {
action = Action::Right;
} else if cursor.goto_parent() {
action = Action::Up;
} else {
break;
}
}
}
}
let _ = extract_module(parse_tree.root_node(), src.as_bytes());
// println!(
// "{:#?}",
// extract_module(parse_tree.root_node(), src.as_bytes())
// );
}
Ok(())