codebase_to_prompt/
lib.rs

1use std::fs::{self, File};
2use std::io::{self, BufWriter, Write};
3use std::path::{Path, PathBuf};
4
5use anyhow::{Context, Result};
6use chrono::Local;
7use clap::ValueEnum;
8use git2::Repository;
9use ignore::gitignore::Gitignore;
10use tracing::{error, info, warn};
11use walkdir::{DirEntry, WalkDir};
12
13/// Represents the output format for the bundled files.
14///
15/// - `Markdown`: Outputs files in Markdown format with code blocks.
16/// - `Text`: Outputs files as plain text.
17/// - `Console`: Outputs files formatted for console display (default).
18#[derive(Debug, Clone, ValueEnum, Default)]
19pub enum Format {
20    Markdown,
21    Text,
22    #[default]
23    Console,
24}
25
26/// Configuration options for the file bundling process.
27#[derive(Debug)]
28pub struct Config {
29    /// The directory to process.
30    pub directory: PathBuf,
31    /// The optional output file path. If not provided, output is written to stdout.
32    pub output: Option<PathBuf>,
33    /// File extensions to include in the output.
34    pub include: Vec<String>,
35    /// File extensions to exclude from the output.
36    pub exclude: Vec<String>,
37    /// The format of the output (Markdown, Text, or Console).
38    pub format: Format,
39    /// Whether to append the current date to the output file name.
40    pub append_date: bool,
41    /// Whether to append the current Git hash to the output file name.
42    pub append_git_hash: bool,
43    /// Whether to include line numbers in the output.
44    pub line_numbers: bool,
45    /// Whether to ignore hidden files and directories.
46    pub ignore_hidden: bool,
47    /// Whether to respect `.gitignore` rules.
48    pub respect_gitignore: bool,
49}
50
51/// Runs the file bundling process based on the provided configuration.
52///
53/// # Arguments
54/// * `config` - The configuration options for the bundling process.
55///
56/// # Returns
57/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the process fails.
58pub fn run(config: Config) -> Result<()> {
59    let mut output_path = config.output.clone();
60
61    if config.append_date || config.append_git_hash {
62        append_date_and_git_hash(&mut output_path, &config)?;
63    }
64
65    let writer = determine_output_writer(&output_path)?;
66
67    process_directory(&config, writer)
68}
69
70/// Appends the current date and/or Git hash to the output file name if required.
71///
72/// # Arguments
73/// * `output_path` - The optional output file path to modify.
74/// * `config` - The configuration options for the bundling process.
75///
76/// # Returns
77/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the operation fails.
78fn append_date_and_git_hash(output_path: &mut Option<PathBuf>, config: &Config) -> Result<()> {
79    if let Some(path) = output_path {
80        let mut new_filename = path
81            .file_stem()
82            .and_then(|s| s.to_str())
83            .unwrap_or_default()
84            .to_string();
85
86        if config.append_date {
87            new_filename.push('_');
88            new_filename.push_str(&Local::now().format("%Y%m%d").to_string());
89            info!("Appending date to filename.");
90        }
91
92        if config.append_git_hash {
93            match Repository::open(&config.directory) {
94                Ok(repo) => {
95                    let head = repo.head().context("Failed to get repository HEAD")?;
96                    if let Some(oid) = head.target() {
97                        new_filename.push('_');
98                        new_filename.push_str(&oid.to_string()[..7]);
99                        info!("Appending git hash to filename.");
100                    }
101                }
102                Err(_) => warn!("Not a git repository, cannot append git hash."),
103            }
104        }
105
106        if let Some(ext) = path.extension().and_then(|s| s.to_str()) {
107            new_filename.push('.');
108            new_filename.push_str(ext);
109        }
110        path.set_file_name(new_filename);
111    }
112    Ok(())
113}
114
115/// Determines the output writer (file or stdout) based on the configuration.
116///
117/// # Arguments
118/// * `output_path` - The optional output file path.
119///
120/// # Returns
121/// * `Result<Box<dyn Write>>` - Returns a writer for the output.
122fn determine_output_writer(output_path: &Option<PathBuf>) -> Result<Box<dyn Write>> {
123    if let Some(path) = output_path {
124        info!("Output will be written to: {}", path.display());
125        let file = File::create(path)
126            .with_context(|| format!("Failed to create output file: {}", path.display()))?;
127        Ok(Box::new(BufWriter::new(file)))
128    } else {
129        info!("Output will be written to stdout.");
130        Ok(Box::new(BufWriter::new(io::stdout())))
131    }
132}
133
134/// Processes the specified directory and writes the bundled content to the writer.
135///
136/// # Arguments
137/// * `config` - The configuration options for the bundling process.
138/// * `writer` - The writer to output the bundled content.
139///
140/// # Returns
141/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the process fails.
142fn process_directory(config: &Config, mut writer: Box<dyn Write>) -> Result<()> {
143    let (gitignore, _) = Gitignore::new(config.directory.join(".gitignore"));
144
145    let walker = WalkDir::new(&config.directory)
146        .into_iter()
147        .filter_entry(|e| should_include_entry(e, &gitignore, config));
148
149    for result in walker {
150        let entry = match result {
151            Ok(entry) => entry,
152            Err(err) => {
153                error!("Failed to access entry: {}", err);
154                continue;
155            }
156        };
157
158        if let Err(err) = process_file_entry(&entry, &mut writer, config) {
159            error!("{}", err);
160        }
161    }
162
163    info!("File bundling complete.");
164    Ok(())
165}
166
167/// Determines if a directory entry should be included based on the configuration.
168///
169/// # Arguments
170/// * `entry` - The directory entry to check.
171/// * `gitignore` - The `.gitignore` rules to respect.
172/// * `config` - The configuration options for the bundling process.
173///
174/// # Returns
175/// * `bool` - Returns `true` if the entry should be included, `false` otherwise.
176fn should_include_entry(entry: &DirEntry, gitignore: &Gitignore, config: &Config) -> bool {
177    !is_hidden(entry, config) && !is_ignored(entry, gitignore, config)
178}
179
180/// Processes a single file entry and writes its content to the writer.
181///
182/// # Arguments
183/// * `entry` - The file entry to process.
184/// * `writer` - The writer to output the file content.
185/// * `config` - The configuration options for the bundling process.
186///
187/// # Returns
188/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the process fails.
189fn process_file_entry(entry: &DirEntry, writer: &mut dyn Write, config: &Config) -> Result<()> {
190    let path = entry.path();
191    if !path.is_file() {
192        return Ok(());
193    }
194
195    let extension = path.extension().and_then(|s| s.to_str()).unwrap_or("");
196
197    let apply_include_filter =
198        !(config.include.is_empty() || config.include.len() == 1 && config.include[0].is_empty());
199
200    if apply_include_filter && !config.include.contains(&extension.to_string()) {
201        return Ok(());
202    }
203
204    if config.exclude.contains(&extension.to_string()) {
205        return Ok(());
206    }
207
208    let relative_path = path.strip_prefix(&config.directory).unwrap_or(path);
209    let content = match fs::read_to_string(path) {
210        Ok(content) => content,
211        Err(_) => {
212            warn!("Skipping non-UTF-8 file: {}", path.display());
213            return Ok(()); // Skip non-text files
214        }
215    };
216
217    write_file_content(writer, relative_path, &content, extension, config)
218        .with_context(|| format!("Failed to write file content for {}", path.display()))
219}
220
221/// Writes the content of a single file to the writer based on the specified format.
222///
223/// # Arguments
224/// * `writer` - The writer to output the file content.
225/// * `path` - The relative path of the file.
226/// * `content` - The content of the file.
227/// * `extension` - The file extension.
228/// * `config` - The configuration options for the bundling process.
229///
230/// # Returns
231/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the operation fails.
232fn write_file_content(
233    writer: &mut dyn Write,
234    path: &Path,
235    content: &str,
236    extension: &str,
237    config: &Config,
238) -> Result<()> {
239    match config.format {
240        Format::Markdown => {
241            writeln!(writer, "### `{}`\n", path.display())?;
242            writeln!(writer, "```{}", extension)?;
243            write_content_lines(writer, content, config.line_numbers)?;
244            writeln!(writer, "```\n")?;
245        }
246        Format::Text | Format::Console => {
247            // In Console mode, we could add colors or other specific formatting later
248            writeln!(writer, "./{}\n---", path.display())?;
249            write_content_lines(writer, content, config.line_numbers)?;
250            writeln!(writer, "---")?;
251        }
252    }
253    Ok(())
254}
255
256/// Writes content line by line to the writer, optionally including line numbers.
257///
258/// # Arguments
259/// * `writer` - The writer to output the content.
260/// * `content` - The content to write.
261/// * `line_numbers` - Whether to include line numbers.
262///
263/// # Returns
264/// * `Result<()>` - Returns `Ok(())` if successful, or an error if the operation fails.
265fn write_content_lines(writer: &mut dyn Write, content: &str, line_numbers: bool) -> Result<()> {
266    if line_numbers {
267        for (i, line) in content.lines().enumerate() {
268            writeln!(writer, "{:4} | {}", i + 1, line)?;
269        }
270    } else {
271        writeln!(writer, "{}", content)?;
272    }
273    Ok(())
274}
275
276/// Checks if a directory entry is hidden based on the configuration.
277///
278/// # Arguments
279/// * `entry` - The directory entry to check.
280/// * `config` - The configuration options for the bundling process.
281///
282/// # Returns
283/// * `bool` - Returns `true` if the entry is hidden, `false` otherwise.
284fn is_hidden(entry: &DirEntry, config: &Config) -> bool {
285    config.ignore_hidden
286        && entry
287            .file_name()
288            .to_str()
289            .map(|s| s.starts_with('.'))
290            .unwrap_or(false)
291}
292
293/// Checks if a directory entry is ignored by `.gitignore` rules.
294///
295/// # Arguments
296/// * `entry` - The directory entry to check.
297/// * `gitignore` - The `.gitignore` rules to respect.
298/// * `config` - The configuration options for the bundling process.
299///
300/// # Returns
301/// * `bool` - Returns `true` if the entry is ignored, `false` otherwise.
302fn is_ignored(entry: &DirEntry, gitignore: &Gitignore, config: &Config) -> bool {
303    if !config.respect_gitignore {
304        return false;
305    }
306    gitignore
307        .matched(entry.path(), entry.file_type().is_dir())
308        .is_ignore()
309}
codebase_to_prompt/lib.rs

codebase_to_prompt/
lib.rs
