tokenize_cli/
context_generator.rs

1use std::{
2    fs::File,
3    io::{self, BufWriter, Write},
4    sync::{Arc, Mutex},
5};
6
7use ignore::{WalkParallel, WalkState};
8use log::warn;
9
10use crate::{config::Config, file_data::FileData};
11
12pub struct ContextGenerator<'a> {
13    config: &'a Config,
14    writer: Arc<Mutex<BufWriter<File>>>,
15}
16
17impl<'a> ContextGenerator<'a> {
18    pub fn new(config: &'a Config) -> io::Result<Self> {
19        if config.out_file.exists() {
20            warn!(
21                "out file {} already exists and will be overwritten",
22                config.out_file.display()
23            );
24        }
25
26        let out_file = File::create(config.out_file.clone())?;
27        let writer = Arc::new(Mutex::new(BufWriter::new(out_file)));
28
29        Ok(ContextGenerator { config, writer })
30    }
31}
32
33impl ContextGenerator<'_> {
34    pub fn generate(&mut self, walker: WalkParallel) -> io::Result<()> {
35        let default_prompt = include_bytes!("../assets/initial_prompt.md");
36        // let initial_prompt = fs::read(self.config.prompt_file.clone())?;
37        {
38            let mut w = self.writer.lock().unwrap();
39            w.write_all(default_prompt)?;
40        }
41
42        let writer = Arc::clone(&self.writer);
43        let out_filename = self.config.out_file.file_name();
44        walker.run(|| {
45            // This closure runs once per thread, returning the actual visitor
46            let writer = Arc::clone(&writer);
47            Box::new(move |res| {
48                if let Err(err) = &res {
49                    warn!("Walk error: {err}");
50                    return WalkState::Continue;
51                }
52                let entry = res.unwrap();
53
54                if !entry.file_type().is_some_and(|ft| ft.is_file()) {
55                    return WalkState::Continue;
56                }
57                let path = entry.into_path();
58
59                if path.file_name() == out_filename {
60                    return WalkState::Continue;
61                }
62
63                if let Err(e) = (|| -> io::Result<()> {
64                    let fd = FileData::read(&path)?;
65                    fd.write(&mut *writer.lock().unwrap())?;
66                    Ok(())
67                })() {
68                    warn!("{}: {e}", path.display());
69                }
70
71                WalkState::Continue
72            })
73        });
74
75        Ok(())
76    }
77}
tokenize_cli/context_generator.rs

tokenize_cli/
context_generator.rs
