tree_sitter_cli/
input.rs

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