diff --git a/Cargo.lock b/Cargo.lock index dc4f579..ea331f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -850,6 +850,23 @@ dependencies = [ "serde", ] +[[package]] +name = "is-terminal" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e04d7f318608d35d4b61ddd75cbdaee86b023ebe2bd5a66ee0915f0bf93095a9" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.52.0", +] + +[[package]] +name = "is_ci" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" + [[package]] name = "is_terminal_polyfill" version = "1.70.1" @@ -1050,6 +1067,10 @@ name = "owo-colors" version = "4.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48dd4f4a2c8405440fd0462561f0e5806bd0f77e86f51c761481bdd4018b545e" +dependencies = [ + "supports-color 2.1.0", + "supports-color 3.0.2", +] [[package]] name = "parking_lot" @@ -1508,6 +1529,25 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "supports-color" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6398cde53adc3c4557306a96ce67b302968513830a77a95b2b17305d9719a89" +dependencies = [ + "is-terminal", + "is_ci", +] + +[[package]] +name = "supports-color" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c64fc7232dd8d2e4ac5ce4ef302b1d81e0b80d055b9d77c7c4f51f6aa4c867d6" +dependencies = [ + "is_ci", +] + [[package]] name = "syn" version = "2.0.104" @@ -1896,6 +1936,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "supports-color 3.0.2", "tempfile", "thiserror", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 2122bbe..e0fe9b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,7 +39,8 @@ clap = { version = "4.5", features = ["derive"] } tokio-rayon = "2.1.0" rand = "0.8" -owo-colors = "4.1" +owo-colors = { version = "4.1", features = ["supports-colors"] } +supports-color = "3" itertools = "0.14.0" diff --git a/src/main.rs b/src/main.rs index 47aea4c..02f70eb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ use crate::embeddings::Embedding; use anyhow::Result; use clap::Parser; use owo_colors::OwoColorize; +use supports_color::Stream; use rand::prelude::*; use rand::rngs::StdRng; use std::path::Path; @@ -105,6 +106,13 @@ fn sample_random_chunks( } /// Fast semantic code search powered by AI embeddings and turbopuffer +#[derive(clap::ValueEnum, Clone, Debug)] +enum ColorChoice { + Auto, + Always, + Never, +} + #[derive(Parser)] #[command(name = "tg")] #[command(version = "0.1.0")] @@ -173,12 +181,26 @@ struct Cli { /// Show distance scores in output (lower is better) #[arg(long)] scores: bool, + + /// When to use colors: auto, always, never (like ripgrep) + #[arg(long, value_enum, default_value_t = ColorChoice::Auto)] + color: ColorChoice, } #[tokio::main] async fn main() { let cli = Cli::parse(); turbogrep::set_verbose(cli.verbose); + // Determine effective color choice similar to ripgrep + // Priority: --color overrides, then NO_COLOR env disables, then auto based on terminal support + let mut use_color = match cli.color { + ColorChoice::Always => true, + ColorChoice::Never => false, + ColorChoice::Auto => supports_color::on(Stream::Stdout).is_some(), + }; + if std::env::var_os("NO_COLOR").is_some() { + use_color = false; + } if let Err(e) = config::load_or_init_settings().await { eprintln!("<(°!°)> Error loading settings: {}", e); @@ -214,16 +236,19 @@ async fn main() { for chunk in sampled_chunks { if let Some(content) = &chunk.content { - println!( - "{}", - format!( + if use_color { + let path = format!("{}", chunk.path.magenta().bold()); + let start = format!("{}", chunk.start_line.green()); + let end = format!("{}", chunk.end_line.green()); + println!("{}:{}:{}", path, start, end); + } else { + println!( "{path}:{start_line}:{end_line}", path = chunk.path, start_line = chunk.start_line, end_line = chunk.end_line - ) - .bright_cyan() - ); + ); + } println!("{}", content); println!(); // Empty line separator } @@ -271,6 +296,7 @@ async fn main() { cli.max_count, cli.embedding_concurrency, cli.scores, + use_color, ) .await { @@ -288,6 +314,7 @@ async fn main() { cli.max_count, cli.embedding_concurrency, cli.scores, + use_color, ) .await { @@ -304,6 +331,7 @@ async fn main() { cli.max_count, cli.embedding_concurrency, cli.scores, + use_color, ) .await { diff --git a/src/search.rs b/src/search.rs index ae9f854..3075dfe 100644 --- a/src/search.rs +++ b/src/search.rs @@ -1,6 +1,7 @@ use crate::{chunker, embeddings, project, sync, turbopuffer, vprintln}; use anyhow::Result; use embeddings::Embedding; +use owo_colors::OwoColorize; use std::fs::File; use std::io::{BufRead, BufReader}; use std::path::Path; @@ -51,6 +52,7 @@ fn chunks_to_ripgrep_format( chunks: Vec, root_dir: &str, show_scores: bool, + use_color: bool, ) -> String { chunks .into_iter() @@ -71,13 +73,38 @@ fn chunks_to_ripgrep_format( if show_scores { if let Some(distance) = chunk.distance { + if use_color { + format!( + "{}:{}:{}:{}", + format!("{}", relative_path).magenta().bold(), + format!("{}", chunk.start_line).green(), + format!("{:.4}", distance).yellow(), + preview + ) + } else { + format!( + "{}:{}:{:.4}:{}", + relative_path, chunk.start_line, distance, preview + ) + } + } else if use_color { format!( - "{}:{}:{:.4}:{}", - relative_path, chunk.start_line, distance, preview + "{}:{}:{}:{}", + format!("{}", relative_path).magenta().bold(), + format!("{}", chunk.start_line).green(), + "n/a".yellow(), + preview ) } else { format!("{}:{}:n/a:{}", relative_path, chunk.start_line, preview) } + } else if use_color { + format!( + "{}:{}:{}", + format!("{}", relative_path).magenta().bold(), + format!("{}", chunk.start_line).green(), + preview + ) } else { format!("{}:{}:{}", relative_path, chunk.start_line, preview) } @@ -92,6 +119,7 @@ pub async fn search( max_count: usize, embedding_concurrency: Option, show_scores: bool, + use_color: bool, ) -> Result { let (namespace, root_dir) = project::namespace_and_dir(directory) .map_err(|e| SearchError::NamespaceError(e.to_string()))?; @@ -147,6 +175,7 @@ pub async fn search( results_with_content, &root_dir, show_scores, + use_color, )) } @@ -159,6 +188,7 @@ pub async fn speculate_search( max_count: usize, embedding_concurrency: Option, show_scores: bool, + use_color: bool, ) -> Result { loop { let mut search_task = tokio::spawn({ @@ -171,6 +201,7 @@ pub async fn speculate_search( max_count, embedding_concurrency, show_scores, + use_color, ) .await } @@ -275,7 +306,7 @@ mod tests { distance: None, }]; - let result = chunks_to_ripgrep_format(chunks, "/project", false); + let result = chunks_to_ripgrep_format(chunks, "/project", false, false); let expected = "src/main.rs:10:fn main() {"; assert_eq!(result, expected);