use tree_sitter::Node;
pub(crate) fn node_text<'a>(node: Node<'_>, source: &'a [u8]) -> Option<&'a str> {
let start = node.start_byte();
let end = node.end_byte().min(source.len());
std::str::from_utf8(&source[start..end]).ok()
}
pub(crate) fn unwrap_callee(mut node: Node<'_>) -> Node<'_> {
loop {
match node.kind() {
"parenthesized_expression" => {
let mut next = None;
for i in 0..node.named_child_count() {
if let Some(c) = node.named_child(i) {
next = Some(c);
break;
}
}
match next {
Some(n) => node = n,
None => return node,
}
}
"sequence_expression" => {
let last = (0..node.named_child_count())
.rev()
.find_map(|i| node.named_child(i));
match last {
Some(n) => node = n,
None => return node,
}
}
_ => return node,
}
}
}
pub(crate) fn collect_named_args<'a>(args: Node<'a>) -> Vec<Node<'a>> {
let mut out = Vec::new();
let mut cursor = args.walk();
for child in args.children(&mut cursor) {
if child.is_named() {
out.push(child);
}
}
out
}
pub(crate) fn python_kwarg_value<'a>(
args: &[Node<'a>],
name: &str,
source: &[u8],
) -> Option<Node<'a>> {
for a in args {
if a.kind() != "keyword_argument" {
continue;
}
let arg_name = a
.child_by_field_name("name")
.and_then(|n| node_text(n, source));
if arg_name == Some(name) {
return a.child_by_field_name("value");
}
}
None
}
pub(crate) fn python_kwarg_truthy(
args: &[Node<'_>],
name: &str,
source: &[u8],
unknown_default: bool,
) -> bool {
match python_kwarg_value(args, name, source) {
Some(value) => match value.kind() {
"true" => true,
"false" => false,
_ => unknown_default,
},
None => false,
}
}
pub(crate) fn receiver_chain_label(
node: Node<'_>,
source: &[u8],
call_resolver: Option<&dyn Fn(Node<'_>, &[u8]) -> Option<&'static str>>,
) -> String {
match node.kind() {
"member_expression" | "attribute" => {
if let Some(prop) = node
.child_by_field_name("property")
.or_else(|| node.child_by_field_name("attribute"))
{
if let Some(s) = node_text(prop, source) {
return s.to_lowercase();
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
"selector_expression" => {
if let Some(field) = node.child_by_field_name("field") {
if let Some(s) = node_text(field, source) {
return s.to_lowercase();
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
"member_access_expression" => {
if let Some(name) = node.child_by_field_name("name") {
if let Some(s) = node_text(name, source) {
return s.to_lowercase();
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
"call_expression" => {
if let Some(resolver) = call_resolver {
if let Some(label) = resolver(node, source) {
return label.to_string();
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
"await_expression" | "parenthesized_expression" => {
for i in 0..node.named_child_count() {
if let Some(c) = node.named_child(i) {
return receiver_chain_label(c, source, call_resolver);
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
_ => node_text(node, source).unwrap_or("").to_lowercase(),
}
}
pub(crate) fn receiver_chain_label_go(node: Node<'_>, source: &[u8]) -> String {
match node.kind() {
"selector_expression" => {
if let Some(field) = node.child_by_field_name("field") {
if let Some(s) = node_text(field, source) {
return s.to_lowercase();
}
}
node_text(node, source).unwrap_or("").to_lowercase()
}
_ => node_text(node, source).unwrap_or("").to_lowercase(),
}
}
pub(crate) fn node_at_line<'a>(root: Node<'a>, line: u32) -> Option<Node<'a>> {
if line == 0 {
return None;
}
let target_row = line - 1;
let mut current = root;
let in_range = |n: Node<'_>| {
let s = n.start_position().row as u32;
let e = n.end_position().row as u32;
s <= target_row && target_row <= e
};
if !in_range(current) {
return None;
}
loop {
let mut descended = None;
for i in 0..current.named_child_count() {
if let Some(child) = current.named_child(i) {
if in_range(child) {
descended = Some(child);
break;
}
}
}
match descended {
Some(c) => current = c,
None => return Some(current),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ConditionalKind {
If,
Elif,
Else,
}
pub(crate) fn enclosing_python_conditional<'a>(
node: Node<'a>,
) -> Option<(Node<'a>, ConditionalKind)> {
let mut cur = node.parent()?;
loop {
match cur.kind() {
"if_statement" => {
return Some((cur, ConditionalKind::If));
}
"elif_clause" => return Some((cur, ConditionalKind::Elif)),
"else_clause" => return Some((cur, ConditionalKind::Else)),
"function_definition" | "class_definition" | "module" | "lambda" => return None,
_ => {}
}
cur = cur.parent()?;
}
}
pub(crate) fn enclosing_python_function<'a>(node: Node<'a>) -> Option<Node<'a>> {
let mut cur = node.parent()?;
loop {
if cur.kind() == "function_definition" {
return Some(cur);
}
cur = cur.parent()?;
}
}
pub(crate) fn python_function_param_names(func: Node<'_>, source: &[u8]) -> Vec<String> {
let mut names = Vec::new();
let Some(params) = func.child_by_field_name("parameters") else {
return names;
};
let mut cursor = params.walk();
for child in params.children(&mut cursor) {
if !child.is_named() {
continue;
}
let name_node = match child.kind() {
"identifier" => Some(child),
"default_parameter"
| "typed_parameter"
| "typed_default_parameter"
| "list_splat_pattern"
| "dictionary_splat_pattern" => {
child.child_by_field_name("name").or_else(|| {
let mut c2 = child.walk();
let mut found = None;
for sub in child.children(&mut c2) {
if sub.kind() == "identifier" {
found = Some(sub);
break;
}
}
found
})
}
_ => None,
};
if let Some(id) = name_node {
if let Some(text) = node_text(id, source) {
if text != "self" && text != "cls" {
names.push(text.to_string());
}
}
}
}
names
}
pub(crate) fn python_collect_name_idents(expr: Node<'_>, source: &[u8]) -> Vec<String> {
let mut out = Vec::new();
let mut stack = vec![expr];
while let Some(n) = stack.pop() {
match n.kind() {
"identifier" => {
if let Some(t) = node_text(n, source) {
out.push(t.to_string());
}
}
"attribute" => {
if let Some(obj) = n.child_by_field_name("object") {
stack.push(obj);
}
}
_ => {
for i in 0..n.named_child_count() {
if let Some(c) = n.named_child(i) {
stack.push(c);
}
}
}
}
}
out
}
pub(crate) fn python_param_gated_branch_param_name(
node: Node<'_>,
source: &[u8],
) -> Option<String> {
let (clause, kind) = enclosing_python_conditional(node)?;
let func = enclosing_python_function(clause)?;
let condition_node = match kind {
ConditionalKind::If | ConditionalKind::Elif => clause.child_by_field_name("condition")?,
ConditionalKind::Else => {
let parent = clause.parent()?;
if parent.kind() != "if_statement" {
return None;
}
parent.child_by_field_name("condition")?
}
};
let params = python_function_param_names(func, source);
if params.is_empty() {
return None;
}
let names = python_collect_name_idents(condition_node, source);
names.into_iter().find(|n| params.iter().any(|p| p == n))
}
#[cfg(test)]
mod tests {
use super::*;
use tree_sitter::Parser;
fn parse_js(src: &str) -> tree_sitter::Tree {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_javascript::LANGUAGE.into())
.expect("load js grammar");
parser.parse(src, None).expect("parse")
}
fn parse_ts(src: &str) -> tree_sitter::Tree {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into())
.expect("load ts grammar");
parser.parse(src, None).expect("parse")
}
fn parse_python(src: &str) -> tree_sitter::Tree {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_python::LANGUAGE.into())
.expect("load python grammar");
parser.parse(src, None).expect("parse")
}
fn parse_go(src: &str) -> tree_sitter::Tree {
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_go::LANGUAGE.into())
.expect("load go grammar");
parser.parse(src, None).expect("parse")
}
fn find_kind<'a>(root: Node<'a>, kind: &str) -> Option<Node<'a>> {
let mut stack = vec![root];
while let Some(n) = stack.pop() {
if n.kind() == kind {
return Some(n);
}
for i in (0..n.child_count()).rev() {
if let Some(c) = n.child(i) {
stack.push(c);
}
}
}
None
}
#[test]
fn unwrap_callee_parenthesized() {
let src = "(eval)('x');\n";
let tree = parse_js(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let func = call.child_by_field_name("function").expect("func");
assert_eq!(func.kind(), "parenthesized_expression");
let unwrapped = unwrap_callee(func);
assert_eq!(unwrapped.kind(), "identifier");
assert_eq!(node_text(unwrapped, src.as_bytes()), Some("eval"));
}
#[test]
fn unwrap_callee_ts_passthrough() {
let src = "(x as Foo)();\n";
let tree = parse_ts(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let func = call.child_by_field_name("function").expect("func");
let unwrapped = unwrap_callee(func);
assert_eq!(unwrapped.kind(), "as_expression");
}
#[test]
fn unwrap_callee_plain_identifier() {
let src = "eval('x');\n";
let tree = parse_js(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let func = call.child_by_field_name("function").expect("func");
let unwrapped = unwrap_callee(func);
assert_eq!(unwrapped.kind(), "identifier");
assert_eq!(unwrapped.id(), func.id());
}
#[test]
fn collect_named_args_positional_only() {
let src = "f(1, 2, 3);\n";
let tree = parse_js(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let args = call.child_by_field_name("arguments").expect("args");
let collected = collect_named_args(args);
assert_eq!(collected.len(), 3);
assert!(collected.iter().all(|n| n.kind() == "number"));
}
#[test]
fn collect_named_args_python_keyword() {
let src = "f(name='x')\n";
let tree = parse_python(src);
let call = find_kind(tree.root_node(), "call").expect("call");
let args = call.child_by_field_name("arguments").expect("args");
let collected = collect_named_args(args);
assert_eq!(collected.len(), 1);
assert_eq!(collected[0].kind(), "keyword_argument");
}
#[test]
fn collect_named_args_mixed() {
let src = "f(1, 'two', x);\n";
let tree = parse_js(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let args = call.child_by_field_name("arguments").expect("args");
let collected = collect_named_args(args);
assert_eq!(collected.len(), 3);
assert!(!collected.iter().any(|n| n.kind() == ","));
}
fn collect_py_call_args<'a>(tree: &'a tree_sitter::Tree) -> Vec<Node<'a>> {
let call = find_kind(tree.root_node(), "call").expect("call");
let args = call.child_by_field_name("arguments").expect("args");
collect_named_args(args)
}
#[test]
fn python_kwarg_value_locates_named_arg() {
let src = "f(shell=True, timeout=30)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
let v = python_kwarg_value(&args, "shell", src.as_bytes()).expect("shell value");
assert_eq!(v.kind(), "true");
let v2 = python_kwarg_value(&args, "timeout", src.as_bytes()).expect("timeout value");
assert_eq!(v2.kind(), "integer");
}
#[test]
fn python_kwarg_value_returns_none_for_absent() {
let src = "f(shell=True)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(python_kwarg_value(&args, "missing", src.as_bytes()).is_none());
}
#[test]
fn python_kwarg_value_skips_positional_args() {
let src = "f(\"shell\", x)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(python_kwarg_value(&args, "shell", src.as_bytes()).is_none());
}
#[test]
fn python_kwarg_truthy_literal_true() {
let src = "f(shell=True)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(python_kwarg_truthy(&args, "shell", src.as_bytes(), false));
assert!(python_kwarg_truthy(&args, "shell", src.as_bytes(), true));
}
#[test]
fn python_kwarg_truthy_literal_false() {
let src = "f(shell=False)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(!python_kwarg_truthy(&args, "shell", src.as_bytes(), false));
assert!(!python_kwarg_truthy(&args, "shell", src.as_bytes(), true));
}
#[test]
fn python_kwarg_truthy_non_literal_uses_default() {
let src = "f(shell=some_var)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(python_kwarg_truthy(&args, "shell", src.as_bytes(), true));
assert!(!python_kwarg_truthy(&args, "shell", src.as_bytes(), false));
}
#[test]
fn python_kwarg_truthy_absent_is_false() {
let src = "f(timeout=5)\n";
let tree = parse_python(src);
let args = collect_py_call_args(&tree);
assert!(!python_kwarg_truthy(&args, "shell", src.as_bytes(), true));
assert!(!python_kwarg_truthy(&args, "shell", src.as_bytes(), false));
}
#[test]
fn receiver_chain_label_simple_identifier() {
let src = "x;\n";
let tree = parse_js(src);
let id = find_kind(tree.root_node(), "identifier").expect("id");
assert_eq!(receiver_chain_label(id, src.as_bytes(), None), "x");
}
#[test]
fn receiver_chain_label_member_chain() {
let src = "a.b.c;\n";
let tree = parse_js(src);
let mem = find_kind(tree.root_node(), "member_expression").expect("mem");
let label = receiver_chain_label(mem, src.as_bytes(), None);
assert_eq!(label, "c");
}
#[test]
fn receiver_chain_label_python_attribute() {
let src = "os.path.join\n";
let tree = parse_python(src);
let attr = find_kind(tree.root_node(), "attribute").expect("attr");
let label = receiver_chain_label(attr, src.as_bytes(), None);
assert_eq!(label, "join");
}
#[test]
fn receiver_chain_label_call_with_resolver() {
let src = "require('fs');\n";
let tree = parse_js(src);
let call = find_kind(tree.root_node(), "call_expression").expect("call");
let resolver = |n: Node<'_>, _src: &[u8]| -> Option<&'static str> {
if n.kind() == "call_expression" {
Some("fs")
} else {
None
}
};
let label = receiver_chain_label(call, src.as_bytes(), Some(&resolver));
assert_eq!(label, "fs");
let label_noresolve = receiver_chain_label(call, src.as_bytes(), None);
assert_eq!(label_noresolve, "require('fs')");
}
#[test]
fn receiver_chain_label_go_selector() {
let src = "package m\nfunc f() { exec.Command(\"ls\") }\n";
let tree = parse_go(src);
let sel = find_kind(tree.root_node(), "selector_expression").expect("sel");
let label = receiver_chain_label_go(sel, src.as_bytes());
assert_eq!(label, "command");
}
fn node_for_line<'a>(tree: &'a tree_sitter::Tree, line: u32) -> Node<'a> {
node_at_line(tree.root_node(), line).expect("node for line")
}
#[test]
fn param_gated_module_level_returns_false() {
let src = "if condition:\n x = ssl.CERT_NONE\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 2);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn param_gated_if_body_with_matching_param_returns_true() {
let src = "def f(verify):\n if verify is False:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 3);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_global_in_condition_returns_false() {
let src = "def f(x):\n if is_production():\n ctx.verify_mode = ssl.CERT_NONE\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 3);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn param_gated_any_param_match_is_enough() {
let src = "def f(a, verify, c):\n if verify:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 3);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_walks_to_nearest_function() {
let src =
"def outer(verify):\n def inner(other):\n if verify:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 4);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn param_gated_kwargs_count_as_param() {
let src = "def f(**kwargs):\n if kwargs.get('insecure'):\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 3);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_self_does_not_count() {
let src = "class C:\n def m(self):\n if self.disabled:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 4);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn param_gated_cls_does_not_count() {
let src = "class C:\n @classmethod\n def m(cls):\n if cls.disabled:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 5);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn param_gated_else_branch_inherits_if_param_gate() {
let src =
"def f(verify):\n if verify:\n a = 1\n else:\n b = ssl.CERT_NONE\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 5);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_elif_branch_with_param_match() {
let src = "def f(verify):\n if verify is True:\n a = 1\n elif verify is False:\n ctx.verify_mode = ssl.CERT_NONE\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 5);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_typed_default_param() {
let src = "def f(verify: bool = True):\n if verify:\n a = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 3);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_some());
}
#[test]
fn param_gated_function_with_no_useful_params() {
let src = "class C:\n def m(self):\n if some_global:\n x = 1\n";
let tree = parse_python(src);
let node = node_for_line(&tree, 4);
assert!(python_param_gated_branch_param_name(node, src.as_bytes()).is_none());
}
#[test]
fn node_at_line_returns_none_for_out_of_range() {
let src = "x = 1\n";
let tree = parse_python(src);
assert!(node_at_line(tree.root_node(), 999).is_none());
assert!(node_at_line(tree.root_node(), 0).is_none());
}
#[test]
fn node_at_line_finds_deepest_node() {
let src = "def f():\n x = ssl.CERT_NONE\n";
let tree = parse_python(src);
let n = node_at_line(tree.root_node(), 2).expect("node");
assert_ne!(n.kind(), "function_definition");
assert_ne!(n.kind(), "module");
}
}