tree_sitter_cli/
input.rs

1use std::{
2    fs,
3    io::{Read, Write},
4    path::{Path, PathBuf},
5    sync::{
6        atomic::{AtomicUsize, Ordering},
7        mpsc, Arc,
8    },
9};
10
11use anyhow::{anyhow, bail, Context, Result};
12use glob::glob;
13
14use crate::test::{parse_tests, TestEntry};
15
16pub enum CliInput {
17    Paths(Vec<PathBuf>),
18    Test {
19        name: String,
20        contents: Vec<u8>,
21        languages: Vec<Box<str>>,
22    },
23    Stdin(Vec<u8>),
24}
25
26pub fn get_input(
27    paths_file: Option<&Path>,
28    paths: Option<Vec<PathBuf>>,
29    test_number: Option<u32>,
30    cancellation_flag: &Arc<AtomicUsize>,
31) -> Result<CliInput> {
32    if let Some(paths_file) = paths_file {
33        return Ok(CliInput::Paths(
34            fs::read_to_string(paths_file)
35                .with_context(|| format!("Failed to read paths file {}", paths_file.display()))?
36                .trim()
37                .lines()
38                .map(PathBuf::from)
39                .collect::<Vec<_>>(),
40        ));
41    }
42
43    if let Some(test_number) = test_number {
44        let current_dir = std::env::current_dir().unwrap();
45        let test_dir = current_dir.join("test").join("corpus");
46
47        if !test_dir.exists() {
48            return Err(anyhow!(
49                "Test corpus directory not found in current directory, see https://tree-sitter.github.io/tree-sitter/creating-parsers/5-writing-tests"
50            ));
51        }
52
53        let test_entry = parse_tests(&test_dir)?;
54        let mut test_num = 0;
55        let Some((name, contents, languages)) =
56            get_test_info(&test_entry, test_number.max(1) - 1, &mut test_num)
57        else {
58            return Err(anyhow!("Failed to fetch contents of test #{test_number}"));
59        };
60
61        return Ok(CliInput::Test {
62            name,
63            contents,
64            languages,
65        });
66    }
67
68    if let Some(paths) = paths {
69        let mut result = Vec::new();
70
71        let mut incorporate_path = |path: PathBuf, positive| {
72            if positive {
73                result.push(path);
74            } else if let Some(index) = result.iter().position(|p| *p == path) {
75                result.remove(index);
76            }
77        };
78
79        for mut path in paths {
80            let mut positive = true;
81            if path.starts_with("!") {
82                positive = false;
83                path = path.strip_prefix("!").unwrap().to_path_buf();
84            }
85
86            if path.exists() {
87                incorporate_path(path, positive);
88            } else {
89                let Some(path_str) = path.to_str() else {
90                    bail!("Invalid path: {}", path.display());
91                };
92                let paths =
93                    glob(path_str).with_context(|| format!("Invalid glob pattern {path:?}"))?;
94                for path in paths {
95                    incorporate_path(path?, positive);
96                }
97            }
98        }
99
100        if result.is_empty() {
101            return Err(anyhow!(
102                "No files were found at or matched by the provided pathname/glob"
103            ));
104        }
105
106        return Ok(CliInput::Paths(result));
107    }
108
109    let reader_flag = cancellation_flag.clone();
110    let (tx, rx) = mpsc::channel();
111
112    // Spawn a thread to read from stdin, until ctrl-c or EOF is received
113    std::thread::spawn(move || {
114        let mut input = Vec::new();
115        let stdin = std::io::stdin();
116        let mut handle = stdin.lock();
117
118        // Read in chunks, so we can check the ctrl-c flag
119        loop {
120            if reader_flag.load(Ordering::Relaxed) == 1 {
121                break;
122            }
123            let mut buffer = [0; 1024];
124            match handle.read(&mut buffer) {
125                Ok(0) | Err(_) => break,
126                Ok(n) => input.extend_from_slice(&buffer[..n]),
127            }
128        }
129
130        // Signal to the main thread that we're done
131        tx.send(input).ok();
132    });
133
134    loop {
135        // If we've received a ctrl-c signal, exit
136        if cancellation_flag.load(Ordering::Relaxed) == 1 {
137            bail!("\n");
138        }
139
140        // If we're done receiving input from stdin, return it
141        if let Ok(input) = rx.try_recv() {
142            return Ok(CliInput::Stdin(input));
143        }
144
145        std::thread::sleep(std::time::Duration::from_millis(50));
146    }
147}
148
149#[allow(clippy::type_complexity)]
150pub fn get_test_info(
151    test_entry: &TestEntry,
152    target_test: u32,
153    test_num: &mut u32,
154) -> Option<(String, Vec<u8>, Vec<Box<str>>)> {
155    match test_entry {
156        TestEntry::Example {
157            name,
158            input,
159            attributes,
160            ..
161        } => {
162            if *test_num == target_test {
163                return Some((name.clone(), input.clone(), attributes.languages.clone()));
164            }
165            *test_num += 1;
166        }
167        TestEntry::Group { children, .. } => {
168            for child in children {
169                if let Some((name, input, languages)) = get_test_info(child, target_test, test_num)
170                {
171                    return Some((name, input, languages));
172                }
173            }
174        }
175    }
176
177    None
178}
179
180/// Writes `contents` to a temporary file and returns the path to that file.
181pub fn get_tmp_source_file(contents: &[u8]) -> Result<PathBuf> {
182    let parse_path = std::env::temp_dir().join(".tree-sitter-temp");
183    let mut parse_file = std::fs::File::create(&parse_path)?;
184    parse_file.write_all(contents)?;
185
186    Ok(parse_path)
187}
tree_sitter_cli/input.rs

tree_sitter_cli/
input.rs
