From cc449ad9658cb84597f115c51e875cf96bc1cd76 Mon Sep 17 00:00:00 2001 From: Amaan Qureshi Date: Fri, 3 Jan 2025 04:11:37 -0500 Subject: [PATCH] feat(cli): make input handling agnostic Co-authored-by: Will Lillis --- cli/src/input.rs | 187 +++++++++++++++++++++++++++++++++++++++++++++++ cli/src/lib.rs | 1 + cli/src/main.rs | 53 -------------- cli/src/test.rs | 58 --------------- 4 files changed, 188 insertions(+), 111 deletions(-) create mode 100644 cli/src/input.rs diff --git a/cli/src/input.rs b/cli/src/input.rs new file mode 100644 index 00000000..e22d19ef --- /dev/null +++ b/cli/src/input.rs @@ -0,0 +1,187 @@ +use std::{ + fs, + io::{Read, Write}, + path::{Path, PathBuf}, + sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, Arc, + }, +}; + +use anyhow::{anyhow, bail, Context, Result}; +use glob::glob; + +use crate::test::{parse_tests, TestEntry}; + +pub enum CliInput { + Paths(Vec), + Test { + name: String, + contents: Vec, + languages: Vec>, + }, + Stdin(Vec), +} + +pub fn get_input( + paths_file: Option<&Path>, + paths: Option>, + test_number: Option, + cancellation_flag: &Arc, +) -> Result { + if let Some(paths_file) = paths_file { + return Ok(CliInput::Paths( + fs::read_to_string(paths_file) + .with_context(|| format!("Failed to read paths file {}", paths_file.display()))? + .trim() + .lines() + .map(PathBuf::from) + .collect::>(), + )); + } + + if let Some(test_number) = test_number { + let current_dir = std::env::current_dir().unwrap(); + let test_dir = current_dir.join("test").join("corpus"); + + if !test_dir.exists() { + return Err(anyhow!( + "Test corpus directory not found in current directory, see https://tree-sitter.github.io/tree-sitter/creating-parsers/5-writing-tests" + )); + } + + let test_entry = parse_tests(&test_dir)?; + let mut test_num = 0; + let Some((name, contents, languages)) = + get_test_info(&test_entry, test_number.max(1) - 1, &mut test_num) + else { + return Err(anyhow!("Failed to fetch contents of test #{test_number}")); + }; + + return Ok(CliInput::Test { + name, + contents, + languages, + }); + } + + if let Some(paths) = paths { + let mut result = Vec::new(); + + let mut incorporate_path = |path: PathBuf, positive| { + if positive { + result.push(path); + } else if let Some(index) = result.iter().position(|p| *p == path) { + result.remove(index); + } + }; + + for mut path in paths { + let mut positive = true; + if path.starts_with("!") { + positive = false; + path = path.strip_prefix("!").unwrap().to_path_buf(); + } + + if path.exists() { + incorporate_path(path, positive); + } else { + let Some(path_str) = path.to_str() else { + bail!("Invalid path: {}", path.display()); + }; + let paths = + glob(path_str).with_context(|| format!("Invalid glob pattern {path:?}"))?; + for path in paths { + incorporate_path(path?, positive); + } + } + } + + if result.is_empty() { + return Err(anyhow!( + "No files were found at or matched by the provided pathname/glob" + )); + } + + return Ok(CliInput::Paths(result)); + } + + let reader_flag = cancellation_flag.clone(); + let (tx, rx) = mpsc::channel(); + + // Spawn a thread to read from stdin, until ctrl-c or EOF is received + std::thread::spawn(move || { + let mut input = Vec::new(); + let stdin = std::io::stdin(); + let mut handle = stdin.lock(); + + // Read in chunks, so we can check the ctrl-c flag + loop { + if reader_flag.load(Ordering::Relaxed) == 1 { + break; + } + let mut buffer = [0; 1024]; + match handle.read(&mut buffer) { + Ok(0) | Err(_) => break, + Ok(n) => input.extend_from_slice(&buffer[..n]), + } + } + + // Signal to the main thread that we're done + tx.send(input).ok(); + }); + + loop { + // If we've received a ctrl-c signal, exit + if cancellation_flag.load(Ordering::Relaxed) == 1 { + bail!("\n"); + } + + // If we're done receiving input from stdin, return it + if let Ok(input) = rx.try_recv() { + return Ok(CliInput::Stdin(input)); + } + + std::thread::sleep(std::time::Duration::from_millis(50)); + } +} + +#[allow(clippy::type_complexity)] +pub fn get_test_info( + test_entry: &TestEntry, + target_test: u32, + test_num: &mut u32, +) -> Option<(String, Vec, Vec>)> { + match test_entry { + TestEntry::Example { + name, + input, + attributes, + .. + } => { + if *test_num == target_test { + return Some((name.clone(), input.clone(), attributes.languages.clone())); + } + *test_num += 1; + } + TestEntry::Group { children, .. } => { + for child in children { + if let Some((name, input, languages)) = get_test_info(child, target_test, test_num) + { + return Some((name, input, languages)); + } + } + } + } + + None +} + +/// Writes `contents` to a temporary file and returns the path to that file. +pub fn get_tmp_source_file(contents: &[u8]) -> Result { + let parse_path = std::env::temp_dir().join(".tree-sitter-temp"); + let mut parse_file = std::fs::File::create(&parse_path)?; + parse_file.write_all(contents)?; + + Ok(parse_path) +} diff --git a/cli/src/lib.rs b/cli/src/lib.rs index 657e9d9c..9f42a835 100644 --- a/cli/src/lib.rs +++ b/cli/src/lib.rs @@ -3,6 +3,7 @@ pub mod fuzz; pub mod highlight; pub mod init; +pub mod input; pub mod logger; pub mod parse; pub mod playground; diff --git a/cli/src/main.rs b/cli/src/main.rs index 9b85c71d..5fac8d6c 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1446,56 +1446,3 @@ const fn get_styles() -> clap::builder::Styles { ) .placeholder(Style::new().fg_color(Some(Color::Ansi(AnsiColor::White)))) } - -fn collect_paths(paths_file: Option<&str>, paths: Option>) -> Result> { - if let Some(paths_file) = paths_file { - return Ok(fs::read_to_string(paths_file) - .with_context(|| format!("Failed to read paths file {paths_file}"))? - .trim() - .lines() - .map(String::from) - .collect::>()); - } - - if let Some(paths) = paths { - let mut result = Vec::new(); - - let mut incorporate_path = |path: &str, positive| { - if positive { - result.push(path.to_string()); - } else if let Some(index) = result.iter().position(|p| p == path) { - result.remove(index); - } - }; - - for mut path in paths { - let mut positive = true; - if path.starts_with('!') { - positive = false; - path = path.trim_start_matches('!').to_string(); - } - - if Path::new(&path).exists() { - incorporate_path(&path, positive); - } else { - let paths = - glob(&path).with_context(|| format!("Invalid glob pattern {path:?}"))?; - for path in paths { - if let Some(path) = path?.to_str() { - incorporate_path(path, positive); - } - } - } - } - - if result.is_empty() { - return Err(anyhow!( - "No files were found at or matched by the provided pathname/glob" - )); - } - - return Ok(result); - } - - Err(anyhow!("Must provide one or more paths")) -} diff --git a/cli/src/test.rs b/cli/src/test.rs index 11c3bfce..c561405a 100644 --- a/cli/src/test.rs +++ b/cli/src/test.rs @@ -268,64 +268,6 @@ pub fn run_tests_at_path(parser: &mut Parser, opts: &mut TestOptions) -> Result< } } -#[allow(clippy::type_complexity)] -pub fn get_test_info<'test>( - test_entry: &'test TestEntry, - target_test: u32, - test_num: &mut u32, -) -> Option<(&'test str, &'test [u8], Vec>)> { - match test_entry { - TestEntry::Example { - name, - input, - attributes, - .. - } => { - if *test_num == target_test { - return Some((name, input, attributes.languages.clone())); - } - *test_num += 1; - } - TestEntry::Group { children, .. } => { - for child in children { - if let Some((name, input, languages)) = get_test_info(child, target_test, test_num) - { - return Some((name, input, languages)); - } - } - } - } - - None -} - -/// Writes the input of `target_test` to a temporary file and returns the path -pub fn get_tmp_test_file(target_test: u32, color: bool) -> Result<(PathBuf, Vec>)> { - let current_dir = std::env::current_dir().unwrap(); - let test_dir = current_dir.join("test").join("corpus"); - - // Get the input of the target test - let test_entry = parse_tests(&test_dir)?; - let mut test_num = 0; - let Some((test_name, test_contents, languages)) = - get_test_info(&test_entry, target_test - 1, &mut test_num) - else { - return Err(anyhow!("Failed to fetch contents of test #{target_test}")); - }; - - // Write the test contents to a temporary file - let test_path = std::env::temp_dir().join(".tree-sitter-test"); - let mut test_file = std::fs::File::create(&test_path)?; - test_file.write_all(test_contents)?; - - println!( - "{target_test}. {}\n", - paint(color.then_some(AnsiColor::Green), test_name) - ); - - Ok((test_path, languages)) -} - pub fn check_queries_at_path(language: &Language, path: &Path) -> Result<()> { if path.exists() { for entry in WalkDir::new(path)