From cdae6cf1f3a26cbb8482edfd0f80b5f1d6173aab Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Thu, 6 Oct 2022 22:17:39 -0400 Subject: [PATCH] Run first lint rules atop tree-sitter --- examples/parse_python.rs | 468 +------------- foo.py | 1 + src/lib.rs | 19 +- src/linter.rs | 54 +- src/noqa.rs | 1 + src/tree_parser.rs | 1300 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 1341 insertions(+), 502 deletions(-) create mode 100644 foo.py create mode 100644 src/tree_parser.rs diff --git a/examples/parse_python.rs b/examples/parse_python.rs index 2e3e4113c3..ac5f1c9549 100644 --- a/examples/parse_python.rs +++ b/examples/parse_python.rs @@ -1,459 +1,23 @@ -extern crate core; +use std::path::PathBuf; 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}; +use clap::Parser as ClapParser; +use tree_sitter::Parser; -fn to_location(point: Point) -> Location { - Location::new(point.row + 1, point.column + 1) -} +use ruff::fs; +use ruff::tree_parser::extract_module; -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 { - let mut cursor = node.walk(); - node.children(&mut cursor) - .map(|node| extract_statement(node, source)) - .collect() -} - -fn extract_suite(node: Node, source: &[u8]) -> Vec { - 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 { - 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 { - 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, Vec) { - 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::().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()) - } - } +#[derive(Debug, ClapParser)] +struct Cli { + #[arg(required = true)] + file: PathBuf, } fn main() -> Result<()> { - let src = r#" -def double(x): - # Return a double. - return x * 2 + let cli = Cli::parse(); -x = (double(500), double(2, z=1)) -x += 1 + let src = fs::read_file(&cli.file)?; -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()) @@ -461,11 +25,11 @@ baz as wop): let parse_tree = parser.parse(src.as_bytes(), None); if let Some(parse_tree) = &parse_tree { - let _ = extract_module(parse_tree.root_node(), src.as_bytes()); - // println!( - // "{:#?}", - // extract_module(parse_tree.root_node(), src.as_bytes()) - // ); + // let _ = extract_module(parse_tree.root_node(), src.as_bytes()); + println!( + "{:#?}", + extract_module(parse_tree.root_node(), src.as_bytes()) + ); } Ok(()) diff --git a/foo.py b/foo.py new file mode 100644 index 0000000000..7d02ab6a28 --- /dev/null +++ b/foo.py @@ -0,0 +1 @@ +x = call(1) diff --git a/src/lib.rs b/src/lib.rs index 4beb99d634..e6988218b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,10 +2,9 @@ use std::path::Path; use anyhow::Result; use log::debug; -use rustpython_parser::lexer::LexResult; use crate::autofix::fixer::Mode; -use crate::linter::{check_path, tokenize}; +use crate::linter::check_path; use crate::message::Message; use crate::settings::{RawSettings, Settings}; @@ -27,6 +26,7 @@ pub mod printer; pub mod pyproject; mod python; pub mod settings; +pub mod tree_parser; /// Run ruff over Python source code directly. pub fn check(path: &Path, contents: &str) -> Result> { @@ -44,21 +44,8 @@ pub fn check(path: &Path, contents: &str) -> Result> { let settings = Settings::from_raw(RawSettings::from_pyproject(&pyproject, &project_root)?); - // Tokenize once. - let tokens: Vec = tokenize(contents); - - // Determine the noqa line for every line in the source. - let noqa_line_for = noqa::extract_noqa_line_for(&tokens); - // Generate checks. - let checks = check_path( - path, - contents, - tokens, - &noqa_line_for, - &settings, - &Mode::None, - )?; + let checks = check_path(path, contents, &[], &settings, &Mode::None)?; // Convert to messages. let messages: Vec = checks diff --git a/src/linter.rs b/src/linter.rs index bb159a9dce..c42b47a485 100644 --- a/src/linter.rs +++ b/src/linter.rs @@ -5,6 +5,7 @@ use anyhow::Result; use log::debug; use rustpython_parser::lexer::LexResult; use rustpython_parser::{lexer, parser}; +use tree_sitter::Parser; use crate::ast::types::Range; use crate::autofix::fixer; @@ -16,7 +17,8 @@ use crate::code_gen::SourceGenerator; use crate::message::Message; use crate::noqa::add_noqa; use crate::settings::Settings; -use crate::{cache, fs, noqa}; +use crate::tree_parser::extract_module; +use crate::{cache, fs}; /// Collect tokens up to and including the first error. pub(crate) fn tokenize(contents: &str) -> Vec { @@ -34,7 +36,6 @@ pub(crate) fn tokenize(contents: &str) -> Vec { pub(crate) fn check_path( path: &Path, contents: &str, - tokens: Vec, noqa_line_for: &[usize], settings: &Settings, autofix: &fixer::Mode, @@ -48,17 +49,25 @@ pub(crate) fn check_path( .iter() .any(|check_code| matches!(check_code.lint_source(), LintSource::AST)) { - match parser::parse_program_tokens(tokens, "") { + let src = contents.as_bytes(); + + let mut parser = Parser::new(); + parser + .set_language(tree_sitter_python::language()) + .expect("Error loading Python grammar"); + let parse_tree = parser.parse(src, None).unwrap(); + + match extract_module(parse_tree.root_node(), src) { Ok(python_ast) => { checks.extend(check_ast(&python_ast, contents, settings, autofix, path)) } Err(parse_error) => { if settings.enabled.contains(&CheckCode::E999) { checks.push(Check::new( - CheckKind::SyntaxError(parse_error.error.to_string()), + CheckKind::SyntaxError(parse_error.to_string()), Range { - location: parse_error.location, - end_location: parse_error.location, + location: Default::default(), + end_location: Default::default(), }, )) } @@ -100,14 +109,8 @@ pub fn lint_path( // Read the file from disk. let contents = fs::read_file(path)?; - // Tokenize once. - let tokens: Vec = tokenize(&contents); - - // Determine the noqa line for every line in the source. - let noqa_line_for = noqa::extract_noqa_line_for(&tokens); - // Generate checks. - let mut checks = check_path(path, &contents, tokens, &noqa_line_for, settings, autofix)?; + let mut checks = check_path(path, &contents, &[], settings, autofix)?; // Apply autofix. if matches!(autofix, fixer::Mode::Apply) { @@ -134,23 +137,10 @@ pub fn add_noqa_to_path(path: &Path, settings: &Settings) -> Result { // Read the file from disk. let contents = fs::read_file(path)?; - // Tokenize once. - let tokens: Vec = tokenize(&contents); - - // Determine the noqa line for every line in the source. - let noqa_line_for = noqa::extract_noqa_line_for(&tokens); - // Generate checks. - let checks = check_path( - path, - &contents, - tokens, - &noqa_line_for, - settings, - &fixer::Mode::None, - )?; + let checks = check_path(path, &contents, &[], settings, &fixer::Mode::None)?; - add_noqa(&checks, &contents, &noqa_line_for, path) + add_noqa(&checks, &contents, &[], path) } pub fn autoformat_path(path: &Path) -> Result<()> { @@ -175,14 +165,12 @@ mod tests { use anyhow::Result; use regex::Regex; - use rustpython_parser::lexer::LexResult; use crate::autofix::fixer; use crate::checks::{Check, CheckCode}; + use crate::fs; use crate::linter; - use crate::linter::tokenize; use crate::settings; - use crate::{fs, noqa}; fn check_path( path: &Path, @@ -190,9 +178,7 @@ mod tests { autofix: &fixer::Mode, ) -> Result> { let contents = fs::read_file(path)?; - let tokens: Vec = tokenize(&contents); - let noqa_line_for = noqa::extract_noqa_line_for(&tokens); - linter::check_path(path, &contents, tokens, &noqa_line_for, settings, autofix) + linter::check_path(path, &contents, &[], settings, autofix) } #[test] diff --git a/src/noqa.rs b/src/noqa.rs index 8e6abdccf0..8042a4ef7d 100644 --- a/src/noqa.rs +++ b/src/noqa.rs @@ -44,6 +44,7 @@ pub fn extract_noqa_directive(line: &str) -> Directive { } } +#[allow(dead_code)] pub fn extract_noqa_line_for(lxr: &[LexResult]) -> Vec { let mut noqa_line_for: Vec = vec![]; diff --git a/src/tree_parser.rs b/src/tree_parser.rs new file mode 100644 index 0000000000..a7d922209a --- /dev/null +++ b/src/tree_parser.rs @@ -0,0 +1,1300 @@ +use anyhow::Result; +use itertools::any; +use num_bigint::BigInt; +use num_traits::Num; +use rustpython_ast::{ + Alias, AliasData, Arg, ArgData, Arguments, Boolop, Cmpop, Comprehension, Constant, Expr, + ExprContext, ExprKind, Keyword, KeywordData, Location, Operator, Stmt, StmtKind, Unaryop, + Withitem, +}; +use tree_sitter::{Node, Point}; + +#[allow(dead_code)] +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 to_location(point: Point) -> Location { + Location::new(point.row + 1, point.column + 1) +} + +pub fn extract_module(node: Node, source: &[u8]) -> Result> { + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .filter(|node| node.kind() != "comment") + .map(|node| extract_statement(node, source)) + .collect() +} + +fn extract_suite(node: Node, source: &[u8]) -> Result> { + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .filter(|node| node.kind() != "comment") + .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_parameters(node: Node, source: &[u8]) -> Result { + let mut defaults = vec![]; + let mut kw_defaults = vec![]; + let mut kwonlyargs = vec![]; + let mut posonlyargs = vec![]; + let mut args = vec![]; + let mut vararg = None; + let mut kwarg = None; + let mut is_kwonly = false; + for node in node.named_children(&mut node.walk()) { + match node.kind() { + "identifier" => { + let arg = Arg::new( + to_location(node.start_position()), + to_location(node.end_position()), + ArgData { + arg: extract_text(node, source), + annotation: None, + type_comment: None, + }, + ); + if is_kwonly { + kwonlyargs.push(arg) + } else { + args.push(arg) + } + } + "default_parameter" => { + let arg = node.named_child(0).unwrap(); + let default = node.named_child(1).unwrap(); + + let arg = Arg::new( + to_location(arg.start_position()), + to_location(arg.end_position()), + ArgData { + arg: extract_text(arg, source), + annotation: None, + type_comment: None, + }, + ); + let default = extract_expression(default, source)?; + + if is_kwonly { + kwonlyargs.push(arg); + kw_defaults.push(default); + } else { + args.push(arg); + defaults.push(default); + } + } + "typed_parameter" => { + let arg = node.named_child(0).unwrap(); + let _type = node.named_child(1).unwrap(); + + let arg = Arg::new( + to_location(node.start_position()), + to_location(node.end_position()), + ArgData { + arg: extract_text(arg, source), + annotation: Some(Box::new(extract_expression(_type, source)?)), + type_comment: None, + }, + ); + + if is_kwonly { + kwonlyargs.push(arg); + } else { + args.push(arg); + } + } + "typed_default_parameter" => { + let arg = node.named_child(0).unwrap(); + let _type = node.named_child(1).unwrap(); + let default = node.named_child(2).unwrap(); + + let arg = Arg::new( + to_location(node.start_position()), + to_location(node.end_position()), + ArgData { + arg: extract_text(arg, source), + annotation: Some(Box::new(extract_expression(_type, source)?)), + type_comment: None, + }, + ); + let default = extract_expression(default, source)?; + + if is_kwonly { + kwonlyargs.push(arg); + kw_defaults.push(default); + } else { + args.push(arg); + defaults.push(default); + } + } + "positional_separator" => { + // Shift the positional arguments over to positional-only. + while let Some(arg) = args.pop() { + posonlyargs.push(arg); + } + posonlyargs.reverse(); + } + "keyword_separator" => { + is_kwonly = true; + } + "list_splat_pattern" => { + let arg = node.named_child(0).unwrap(); + let arg = Arg::new( + to_location(arg.start_position()), + to_location(arg.end_position()), + ArgData { + arg: extract_text(arg, source), + annotation: None, + type_comment: None, + }, + ); + vararg = Some(Box::new(arg)); + } + "dictionary_splat_pattern" => { + let arg = node.named_child(0).unwrap(); + let arg = Arg::new( + to_location(arg.start_position()), + to_location(arg.end_position()), + ArgData { + arg: extract_text(arg, source), + annotation: None, + type_comment: None, + }, + ); + kwarg = Some(Box::new(arg)); + } + kind => { + return Err(anyhow::anyhow!("Unexpected parameter kind: {}.", kind)); + } + } + } + + Ok(Arguments { + posonlyargs, + args, + vararg, + kwonlyargs, + kw_defaults, + kwarg, + defaults, + }) +} + +fn extract_with_clause(_node: Node, _source: &[u8]) -> Vec { + vec![] +} + +fn extract_import_list(node: Node, source: &[u8]) -> Vec { + let mut aliases = vec![]; + for node in node.children_by_field_name("name", &mut node.walk()) { + // Alias. + if let Some(asname) = node.child_by_field_name("alias") { + let name = node.child_by_field_name("name").unwrap(); + aliases.push(Alias::new( + to_location(node.start_position()), + to_location(node.end_position()), + AliasData { + name: extract_text(name, source), + asname: Some(extract_text(asname, source)), + }, + )); + } else { + let name = node.named_child(0).unwrap(); + aliases.push(Alias::new( + to_location(node.start_position()), + to_location(node.end_position()), + AliasData { + name: extract_text(name, source), + asname: None, + }, + )); + } + } + aliases +} + +fn extract_statement(node: Node, source: &[u8]) -> Result { + match node.kind() { + "for_statement" => Ok(if node.child(0).unwrap().kind() == "async" { + Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::AsyncFor { + 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)?, + orelse: node + .child_by_field_name("alternative") + .map(|node| { + extract_suite(node.child_by_field_name("body").unwrap(), source) + }) + .transpose()? + .unwrap_or_default(), + type_comment: None, + }, + ) + } else { + 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)?, + orelse: node + .child_by_field_name("alternative") + .map(|node| { + extract_suite(node.child_by_field_name("body").unwrap(), source) + }) + .transpose()? + .unwrap_or_default(), + type_comment: None, + }, + ) + }), + "while_statement" => Ok(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)?, + orelse: node + .child_by_field_name("alternative") + .map(|node| extract_suite(node.child_by_field_name("body").unwrap(), source)) + .transpose()? + .unwrap_or_default(), + }, + )), + "with_statement" => Ok(if node.child(0).unwrap().kind() == "async" { + Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::AsyncWith { + items: extract_with_clause(node.named_child(1).unwrap(), source), + body: extract_suite(node.child_by_field_name("body").unwrap(), source)?, + type_comment: None, + }, + ) + } else { + Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::With { + items: extract_with_clause(node.named_child(1).unwrap(), source), + body: extract_suite(node.child_by_field_name("body").unwrap(), source)?, + type_comment: None, + }, + ) + }), + "if_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::If { + test: Box::new(extract_expression( + node.child_by_field_name("condition").unwrap(), + source, + )?), + body: extract_suite(node.child_by_field_name("consequence").unwrap(), source)?, + // TODO(charlie): Unimplemented. + orelse: vec![], + }, + )), + "class_definition" => { + let (bases, keywords) = node + .child_by_field_name("superclasses") + .map(|node| extract_argument_list(node, source)) + .transpose()? + .unwrap_or_default(); + + Ok(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): Unimplemented. + decorator_list: vec![], + }, + )) + } + "function_definition" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::FunctionDef { + name: extract_text(node.named_child(0).unwrap(), source), + args: Box::new(extract_parameters(node.named_child(1).unwrap(), source)?), + body: extract_suite(node.child_by_field_name("body").unwrap(), source)?, + // TODO(charlie): Unimplemented. + decorator_list: vec![], + returns: None, + type_comment: None, + }, + )), + "return_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Return { + value: node + .child(1) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + }, + )), + "pass_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Pass, + )), + "continue_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Continue, + )), + "break_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Break, + )), + "import_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Import { + names: extract_import_list(node, source), + }, + )), + "import_from_statement" | "future_import_statement" => { + let mut cursor = node.walk(); + + // Find the module name. + let module: Option; + let level: Option; + let child = node.named_child(0).unwrap(); + match child.kind() { + "relative_import" => { + level = Some(extract_text(child.named_child(0).unwrap(), source).len()); + module = child.named_child(1).map(|node| extract_text(node, source)); + } + "dotted_name" => { + level = None; + module = Some(extract_text(child, source)); + } + kind => { + return Err(anyhow::anyhow!( + "Expected relative_import or dotted_name; got: {}.", + kind + )); + } + } + + // Find the imports. + let mut names: Vec = vec![]; + for node in node.named_children(&mut cursor).skip(1) { + match node.kind() { + "wildcard_import" => names.push(Alias::new( + to_location(node.start_position()), + to_location(node.end_position()), + AliasData { + name: "*".to_string(), + asname: None, + }, + )), + "aliased_import" => names.push(Alias::new( + to_location(node.start_position()), + to_location(node.end_position()), + AliasData { + name: extract_text(node.child_by_field_name("name").unwrap(), source), + asname: Some(extract_text( + node.child_by_field_name("alias").unwrap(), + source, + )), + }, + )), + "dotted_name" => names.push(Alias::new( + to_location(node.start_position()), + to_location(node.end_position()), + AliasData { + name: extract_text(node, source), + asname: None, + }, + )), + kind => { + return Err(anyhow::anyhow!( + "Expected relative_import or dotted_name; got: {}.", + kind + )); + } + } + } + + Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::ImportFrom { + module, + names, + level, + }, + )) + } + "expression_statement" => { + let node = node.named_child(0).unwrap(); + match node.kind() { + "assignment" => Ok(if let Some(_type) = node.child_by_field_name("type") { + Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::AnnAssign { + target: Box::new(extract_expression( + node.child_by_field_name("left").unwrap(), + source, + )?), + annotation: Box::new(extract_expression(_type, source)?), + value: node + .child_by_field_name("right") + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + // TODO(charlie): Unimplemented. + simple: 0, + }, + ) + } else { + let mut targets = vec![extract_expression( + node.child_by_field_name("left").unwrap(), + source, + )?]; + targets.extend(extract_expression_or_list( + node.child_by_field_name("right").unwrap(), + source, + )?); + let value = Box::new(targets.pop().unwrap()); + Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Assign { + targets, + value, + type_comment: None, + }, + ) + }), + "augmented_assignment" => Ok(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, + ), + }, + )), + _ => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Expr { + value: Box::new(extract_expression(node, source)?), + }, + )), + } + } + "try_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Try { + body: extract_suite(node.child_by_field_name("body").unwrap(), source)?, + // TODO(charlie): Unimplemented. + handlers: vec![], + orelse: vec![], + finalbody: vec![], + }, + )), + "raise_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Raise { + exc: node + .named_child(0) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + cause: node + .named_child(1) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + }, + )), + "decorated_definition" => { + extract_statement(node.child_by_field_name("definition").unwrap(), source) + } + "global_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Global { + names: node + .named_children(&mut node.walk()) + .map(|node| extract_text(node, source)) + .collect(), + }, + )), + "nonlocal_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Nonlocal { + names: node + .named_children(&mut node.walk()) + .map(|node| extract_text(node, source)) + .collect(), + }, + )), + "delete_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Delete { + targets: node + .named_children(&mut node.walk()) + .map(|node| extract_expression(node, source)) + .collect::>>()?, + }, + )), + "assert_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Assert { + test: Box::new(extract_expression(node.named_child(0).unwrap(), source)?), + msg: node + .named_child(1) + .map(|node| Box::new(extract_expression(node, source).unwrap())), + }, + )), + "print_statement" => Ok(Stmt::new( + to_location(node.start_position()), + to_location(node.end_position()), + StmtKind::Expr { + value: Box::new(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Call { + func: Box::new(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Name { + id: "print".to_string(), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + // TODO(charlie): Unimplemented. + args: vec![], + keywords: vec![], + }, + )), + }, + )), + kind => Err(anyhow::anyhow!("Unhandled statement kind: {}", kind)), + } +} + +fn extract_expression_list(node: Node, source: &[u8]) -> Result> { + let mut cursor = node.walk(); + node.named_children(&mut cursor) + .filter(|node| node.kind() != "comment") + .map(|node| extract_expression(node, source)) + .collect() +} + +fn extract_expression_or_list(node: Node, source: &[u8]) -> Result> { + match node.kind() { + "expression_list" + | "pattern_list" + | "tuple_pattern" + | "named_expression" + | "integer" + | "float" + | "concatenated_string" + | "string" + | "tuple" + | "identifier" + | "call" + | "generator_expression" + | "binary_operator" + | "unary_operator" + | "true" + | "false" + | "none" + | "not_operator" + | "boolean_operator" + | "comparison_operator" + | "ellipsis" + | "yield" + | "await" + | "list_comprehension" + | "set_comprehension" + | "dictionary_comprehension" + | "list" + | "set" + | "list_pattern" + | "list_splat" + | "list_splat_pattern" + | "dictionary" + | "type" + | "subscript" + | "attribute" + | "lambda" + | "slice" + | "parenthesized_expression" + | "conditional_expression" => Ok(vec![extract_expression(node, source)?]), + _ => extract_expression_list(node, source), + } +} + +fn extract_keyword_argument(node: Node, source: &[u8]) -> Result { + Ok(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]) -> Result<(Vec, Vec)> { + let mut args = vec![]; + let mut keywords = vec![]; + for child in node.named_children(&mut node.walk()) { + match child.kind() { + "keyword_argument" => { + keywords.push(extract_keyword_argument(child, source)?); + } + _ => args.push(extract_expression(child, source)?), + } + } + Ok((args, keywords)) +} + +fn extract_pair(node: Node, source: &[u8]) -> Result<(Expr, Expr)> { + Ok(( + extract_expression(node.child_by_field_name("key").unwrap(), source)?, + extract_expression(node.child_by_field_name("value").unwrap(), source)?, + )) +} + +fn extract_generators(node: Node, source: &[u8]) -> Result> { + let mut generators: Vec = vec![]; + for node in node.named_children(&mut node.walk()).skip(1) { + if node.child_by_field_name("left").is_some() { + generators.push(Comprehension { + 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, + )?), + is_async: if node.child(0).unwrap().kind() == "async" { + 1 + } else { + 0 + }, + ifs: vec![], + }); + } else { + generators + .last_mut() + .unwrap() + .ifs + .push(extract_expression(node.named_child(0).unwrap(), source)?); + } + } + Ok(generators) +} + +fn extract_expression(node: Node, source: &[u8]) -> Result { + match node.kind() { + "expression_list" | "pattern_list" | "tuple_pattern" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Tuple { + elts: extract_expression_list(node, source)?, + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "named_expression" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::NamedExpr { + target: Box::new(extract_expression( + node.child_by_field_name("name").unwrap(), + source, + )?), + value: Box::new(extract_expression( + node.child_by_field_name("value").unwrap(), + source, + )?), + }, + )), + "integer" => { + let text = extract_text(node, source); + for (pattern, radix) in [ + ("0x", 16), + ("0X", 16), + ("0o", 8), + ("0O", 8), + ("0b", 2), + ("0B", 2), + ] { + if let Some(remainder) = text.strip_prefix(pattern) { + return Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Int(BigInt::from_str_radix(remainder, radix).unwrap()), + kind: None, + }, + )); + } + } + + for pattern in ['j', 'J'] { + if let Some(remainder) = text.strip_suffix(pattern) { + return Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Complex { + real: 0., + imag: remainder.parse::().unwrap(), + }, + kind: None, + }, + )); + } + } + + Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Int(BigInt::from_str_radix(&text, 10).unwrap()), + kind: None, + }, + )) + } + "float" => { + let text = extract_text(node, source); + + for pattern in ['j', 'J'] { + if let Some(remainder) = text.strip_suffix(pattern) { + return Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Complex { + real: 0., + imag: remainder.parse::().unwrap(), + }, + kind: None, + }, + )); + } + } + + Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Float(text.parse::().unwrap()), + kind: None, + }, + )) + } + "concatenated_string" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::JoinedStr { + values: extract_expression_list(node, source)?, + }, + )), + "string" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Str(extract_text(node, source)), + // TODO(charlie): Unimplemented. + kind: None, + }, + )), + "tuple" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Tuple { + elts: extract_expression_list(node, source)?, + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "identifier" => Ok(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(), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "call" => { + let argument_list = + extract_argument_list(node.child_by_field_name("arguments").unwrap(), source)?; + Ok(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, + }, + )) + } + "generator_expression" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::GeneratorExp { + elt: Box::new(extract_expression( + node.child_by_field_name("body").unwrap(), + source, + )?), + generators: extract_generators(node, source)?, + }, + )), + "binary_operator" => Ok(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, + )?), + }, + )), + "unary_operator" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::UnaryOp { + op: match node.child_by_field_name("operator").unwrap().kind() { + "+" => Unaryop::UAdd, + "-" => Unaryop::USub, + "~" => Unaryop::Invert, + op => panic!("Invalid unary operator: {}", op), + }, + operand: Box::new(extract_expression( + node.child_by_field_name("argument").unwrap(), + source, + )?), + }, + )), + "true" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Bool(true), + kind: None, + }, + )), + "false" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Bool(false), + kind: None, + }, + )), + "none" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::None, + kind: None, + }, + )), + "not_operator" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::UnaryOp { + op: Unaryop::Not, + operand: Box::new(extract_expression( + node.child_by_field_name("argument").unwrap(), + source, + )?), + }, + )), + "boolean_operator" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::BoolOp { + op: match node.child_by_field_name("operator").unwrap().kind() { + "and" => Boolop::And, + "or" => Boolop::Or, + op => panic!("Invalid boolean operator: {}", op), + }, + values: vec![ + extract_expression(node.child_by_field_name("left").unwrap(), source)?, + extract_expression(node.child_by_field_name("right").unwrap(), source)?, + ], + }, + )), + "comparison_operator" => { + let mut cursor = node.walk(); + + // Find the left name. + let left = Box::new(extract_expression(node.named_child(0).unwrap(), source)?); + + // Find the comparators. + let ops: Vec = node + .children_by_field_name("operators", &mut cursor) + .map(|node| match node.kind() { + ">" => Ok(Cmpop::Gt), + "<" => Ok(Cmpop::Lt), + ">=" => Ok(Cmpop::GtE), + "<=" => Ok(Cmpop::LtE), + kind => Err(anyhow::anyhow!("Unhandled operator kind: {}", kind)), + }) + .collect::>>()?; + + // Find the other operators. + let comparators: Vec = node + .named_children(&mut cursor) + .skip(1) + .map(|node| extract_expression(node, source)) + .collect::>>()?; + + Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Compare { + left, + ops, + comparators, + }, + )) + } + "ellipsis" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Constant { + value: Constant::Ellipsis, + kind: None, + }, + )), + "yield" => match node.named_child(1) { + None => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Yield { value: None }, + )), + Some(node) => match node.kind() { + "from" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::YieldFrom { + value: Box::new(extract_expression(node.next_sibling().unwrap(), source)?), + }, + )), + _ => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Yield { + value: Some(Box::new(extract_expression(node, source)?)), + }, + )), + }, + }, + "await" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Await { + value: Box::new(extract_expression(node.named_child(0).unwrap(), source)?), + }, + )), + "list_comprehension" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::ListComp { + elt: Box::new(extract_expression( + node.child_by_field_name("body").unwrap(), + source, + )?), + generators: extract_generators(node, source)?, + }, + )), + "set_comprehension" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::SetComp { + elt: Box::new(extract_expression( + node.child_by_field_name("body").unwrap(), + source, + )?), + generators: extract_generators(node, source)?, + }, + )), + "dictionary_comprehension" => { + let (key, value) = extract_pair(node.child_by_field_name("body").unwrap(), source)?; + Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::DictComp { + key: Box::new(key), + value: Box::new(value), + generators: extract_generators(node, source)?, + }, + )) + } + "list" | "set" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::List { + elts: node + .named_children(&mut node.walk()) + .map(|node| extract_expression(node, source)) + .collect::>>()?, + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "list_pattern" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::List { + elts: node + .named_children(&mut node.walk()) + .map(|node| extract_expression(node, source)) + .collect::>>()?, + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "list_splat" | "list_splat_pattern" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Starred { + value: Box::new(extract_expression(node.named_child(0).unwrap(), source)?), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "dictionary" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Dict { + // TODO(charlie): Unimplemented. + keys: vec![], + values: vec![], + }, + )), + "type" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Name { + id: extract_text(node, source), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "subscript" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Subscript { + value: Box::new(extract_expression( + node.child_by_field_name("value").unwrap(), + source, + )?), + slice: Box::new(extract_expression( + node.child_by_field_name("subscript").unwrap(), + source, + )?), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "attribute" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Attribute { + value: Box::new(extract_expression( + node.child_by_field_name("object").unwrap(), + source, + )?), + attr: extract_text(node.child_by_field_name("attribute").unwrap(), source), + // TODO(charlie): Track context. + ctx: ExprContext::Load, + }, + )), + "lambda" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Lambda { + args: node + .child_by_field_name("parameters") + .map(|node| extract_parameters(node, source)) + .transpose()? + .map(Box::new) + .unwrap_or_else(|| { + Box::new(Arguments { + posonlyargs: vec![], + args: vec![], + vararg: None, + kwonlyargs: vec![], + kw_defaults: vec![], + kwarg: None, + defaults: vec![], + }) + }), + body: Box::new(extract_expression( + node.child_by_field_name("body").unwrap(), + source, + )?), + }, + )), + "slice" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::Slice { + lower: node + .named_child(0) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + upper: node + .named_child(1) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + step: node + .named_child(2) + .map(|node| extract_expression(node, source)) + .transpose()? + .map(Box::new), + }, + )), + "parenthesized_expression" => { + for child in node.named_children(&mut node.walk()) { + if child.kind() != "comment" { + return extract_expression(child, source); + } + } + Err(anyhow::anyhow!( + "Unable to find expression within parentheses." + )) + } + "conditional_expression" => Ok(Expr::new( + to_location(node.start_position()), + to_location(node.end_position()), + ExprKind::IfExp { + test: Box::new(extract_expression(node.named_child(0).unwrap(), source)?), + body: Box::new(extract_expression(node.named_child(1).unwrap(), source)?), + orelse: Box::new(extract_expression(node.named_child(2).unwrap(), source)?), + }, + )), + kind => Err(anyhow::anyhow!("Unhandled expression kind: {}", kind)), + } +}