atelier/llmc/src/main.rs
2025-04-12 00:02:40 -04:00

104 lines
2.2 KiB
Rust

use atelier::preludes::fs::*;
use atelier::preludes::io::*;
static USAGE: &str = r#"
llmc - Collect code in a directory into a single text output, for feeding into an LLM
USAGE:
llmc [path] [options]...
OPTIONS:
-i, --input-exts <ext,...> The file extensions to include in the output
"#;
fn find_git_root(start: impl AsRef<Path>) -> IoResult<Option<PathBuf>> {
let mut current = Some(start.as_ref().to_path_buf().canonicalize()?);
while let Some(c) = current {
if c.join(".git").exists() {
return Ok(Some(c.to_path_buf().canonicalize()?));
} else {
current = c.parent().map(|o| o.to_path_buf());
}
}
Ok(None)
}
fn walk_dir(path: impl AsRef<Path>, included_exts: &[&str]) -> IoResult<()> {
let mut types = ignore::types::TypesBuilder::new();
for ext in included_exts.iter() {
types.add(&ext[1..], &format!("*{ext}")).unwrap();
}
types.select("all");
let types = types.build().unwrap();
let mut builder = if let Some(root) = find_git_root(path.as_ref())? {
ignore::WalkBuilder::new(root)
} else {
ignore::WalkBuilder::new(path)
};
builder.types(types);
let walk = builder.build();
for result in walk {
let entry = result.unwrap();
let ft = entry.file_type().unwrap();
if ft.is_dir() {
continue;
}
let path = entry.path().as_os_str().to_string_lossy();
println!("##### Contents of {} #####", path);
println!("{}", std::fs::read_to_string(path.as_ref())?);
}
Ok(())
}
fn main() -> std::io::Result<()> {
let mut included_exts = vec![];
let args: Vec<_> = std::env::args().skip(1).collect();
let mut args = args.iter().map(|s| s.as_str());
let mut params = vec![];
let mut fail = false;
while !fail {
let arg = match args.next() {
None => break,
Some(v) => v,
};
if arg.starts_with('-') {
match arg {
"-h" | "--help" => {
println!("{USAGE}");
return Ok(());
},
"-i" | "--include-exts" => {
fail = args
.next()
.map(|a| included_exts.extend(a.split(',')))
.is_none()
},
_ => (),
}
} else {
params.push(arg);
}
}
if fail {
eprintln!("{USAGE}");
return Ok(());
}
let dir = params.get(0).copied().unwrap_or(".");
walk_dir(dir, &included_exts)?;
Ok(())
}