tree_sitter_cli/
query.rs

1use std::{
2    fs,
3    io::{self, Write},
4    ops::Range,
5    path::Path,
6    time::Instant,
7};
8
9use anstyle::AnsiColor;
10use anyhow::{Context, Result};
11use streaming_iterator::StreamingIterator;
12use tree_sitter::{Language, Parser, Point, Query, QueryCursor};
13
14use crate::{
15    query_testing::{self, to_utf8_point},
16    test::paint,
17};
18
19#[allow(clippy::too_many_arguments)]
20pub fn query_file_at_path(
21    language: &Language,
22    path: &Path,
23    name: &str,
24    query_path: &Path,
25    ordered_captures: bool,
26    byte_range: Option<Range<usize>>,
27    point_range: Option<Range<Point>>,
28    should_test: bool,
29    quiet: bool,
30    print_time: bool,
31    stdin: bool,
32) -> Result<()> {
33    let stdout = io::stdout();
34    let mut stdout = stdout.lock();
35
36    let query_source = fs::read_to_string(query_path)
37        .with_context(|| format!("Error reading query file {query_path:?}"))?;
38    let query = Query::new(language, &query_source).with_context(|| "Query compilation failed")?;
39
40    let mut query_cursor = QueryCursor::new();
41    if let Some(range) = byte_range {
42        query_cursor.set_byte_range(range);
43    }
44    if let Some(range) = point_range {
45        query_cursor.set_point_range(range);
46    }
47
48    let mut parser = Parser::new();
49    parser.set_language(language)?;
50
51    let mut results = Vec::new();
52
53    if !should_test && !stdin {
54        writeln!(&mut stdout, "{name}")?;
55    }
56
57    let source_code =
58        fs::read(path).with_context(|| format!("Error reading source file {path:?}"))?;
59    let tree = parser.parse(&source_code, None).unwrap();
60
61    let start = Instant::now();
62    if ordered_captures {
63        let mut captures = query_cursor.captures(&query, tree.root_node(), source_code.as_slice());
64        while let Some((mat, capture_index)) = captures.next() {
65            let capture = mat.captures[*capture_index];
66            let capture_name = &query.capture_names()[capture.index as usize];
67            if !quiet && !should_test {
68                writeln!(
69                        &mut stdout,
70                        "    pattern: {:>2}, capture: {} - {capture_name}, start: {}, end: {}, text: `{}`",
71                        mat.pattern_index,
72                        capture.index,
73                        capture.node.start_position(),
74                        capture.node.end_position(),
75                        capture.node.utf8_text(&source_code).unwrap_or("")
76                    )?;
77            }
78            results.push(query_testing::CaptureInfo {
79                name: (*capture_name).to_string(),
80                start: to_utf8_point(capture.node.start_position(), source_code.as_slice()),
81                end: to_utf8_point(capture.node.end_position(), source_code.as_slice()),
82            });
83        }
84    } else {
85        let mut matches = query_cursor.matches(&query, tree.root_node(), source_code.as_slice());
86        while let Some(m) = matches.next() {
87            if !quiet && !should_test {
88                writeln!(&mut stdout, "  pattern: {}", m.pattern_index)?;
89            }
90            for capture in m.captures {
91                let start = capture.node.start_position();
92                let end = capture.node.end_position();
93                let capture_name = &query.capture_names()[capture.index as usize];
94                if !quiet && !should_test {
95                    if end.row == start.row {
96                        writeln!(
97                                &mut stdout,
98                                "    capture: {} - {capture_name}, start: {start}, end: {end}, text: `{}`",
99                                capture.index,
100                                capture.node.utf8_text(&source_code).unwrap_or("")
101                            )?;
102                    } else {
103                        writeln!(
104                            &mut stdout,
105                            "    capture: {capture_name}, start: {start}, end: {end}",
106                        )?;
107                    }
108                }
109                results.push(query_testing::CaptureInfo {
110                    name: (*capture_name).to_string(),
111                    start: to_utf8_point(capture.node.start_position(), source_code.as_slice()),
112                    end: to_utf8_point(capture.node.end_position(), source_code.as_slice()),
113                });
114            }
115        }
116    }
117    if query_cursor.did_exceed_match_limit() {
118        writeln!(
119            &mut stdout,
120            "  WARNING: Query exceeded maximum number of in-progress captures!"
121        )?;
122    }
123    if should_test {
124        let path_name = if stdin {
125            "stdin"
126        } else {
127            Path::new(&path).file_name().unwrap().to_str().unwrap()
128        };
129        match query_testing::assert_expected_captures(&results, path, &mut parser, language) {
130            Ok(assertion_count) => {
131                println!(
132                    "  ✓ {} ({} assertions)",
133                    paint(Some(AnsiColor::Green), path_name),
134                    assertion_count
135                );
136            }
137            Err(e) => {
138                println!("  ✗ {}", paint(Some(AnsiColor::Red), path_name));
139                return Err(e);
140            }
141        }
142    }
143    if print_time {
144        writeln!(&mut stdout, "{:?}", start.elapsed())?;
145    }
146
147    Ok(())
148}
tree_sitter_cli/query.rs

tree_sitter_cli/
query.rs
