skubelb/src/rewriter.rs

use anyhow::Result;
use std::collections::HashSet;
use std::path::Path;
use std::{
    fs::{self, File},
    io::{BufReader, prelude::*},
};

pub struct Rewriter {
    source: String,
    replacements: HashSet<String>,
}

impl Rewriter {
    pub fn new(source: String) -> Self {
        Self {
            source,
            replacements: HashSet::new(),
        }
    }

    pub fn add_replacement(&mut self, replacement: String) {
        self.replacements.insert(replacement);
    }

    pub fn remove_replacement(&mut self, replacement: &str) {
        self.replacements.remove(replacement);
    }

    pub fn rewrite_folder(&self, src: &str, dst: &str) -> Result<()> {
        // Make sure we are deterministic; construct a list of strings and sort
        // them
        let mut replacements: Vec<String> = self.replacements.iter().map(|s| s.clone()).collect();
        replacements.sort();
        let dst_base = Path::new(dst);
        let mut to_visit = vec![String::from(src)];
        while to_visit.len() > 0 {
            let dir = to_visit.pop().unwrap();
            for entry in fs::read_dir(&dir)? {
                let entry = entry?;
                let metadata = entry.metadata()?;

                // We need to find the destination this should be written into
                // First, calculate the 'relative' path after we trim the
                // the src prefix, then join that with the dst prefix
                let src_path = entry.path();
                let src_part_path = src_path.strip_prefix(&src)?;
                let dst_path = dst_base.join(src_part_path);

                if metadata.is_dir() {
                    // Create the directory, then carry on. Note that we explore the
                    // src_path after creating dst_path.
                    fs::create_dir(&dst_path)?;
                    to_visit.push(src_path.into_os_string().into_string().unwrap());
                    continue;
                }

                // Open 2 files; one to read and translate, and one to write.
                let source_file = File::open(&src_path)?;
                let mut dest_file = File::create(&dst_path)?;
                let reader = BufReader::new(source_file);

                for line in reader.lines() {
                    let line = line?;
                    // If the line is not subject to replacement, copy it and
                    // carry on.
                    if !line.contains(&self.source) {
                        writeln!(dest_file, "{}", line)?;
                        continue;
                    }

                    // Else, repeat the line multiple times, replacing the string
                    // in question
                    for replacement in &replacements {
                        let new_line = line.replace(&self.source, &replacement);
                        writeln!(dest_file, "{}", new_line)?;
                    }
                }
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use include_directory::{Dir, include_directory};
    use tempdir::TempDir;

    use super::Rewriter;

    static TEST_FILES: Dir<'_> = include_directory!("$CARGO_MANIFEST_DIR/src/testdata");
    #[test]
    fn basic_test() {
        let testdata = TempDir::new("").unwrap();
        TEST_FILES.extract(testdata.path()).unwrap();
        let src = testdata.path().join("testsrc");
        let dst = TempDir::new("").unwrap();

        let mut rewriter = Rewriter::new("to_be_replaced".into());
        rewriter.add_replacement("abc".into());
        rewriter.add_replacement("def".into());
        rewriter.add_replacement("zyx".into());
        rewriter.remove_replacement("zyx");
        rewriter
            .rewrite_folder(
                src.as_os_str().to_str().unwrap(),
                dst.path().as_os_str().to_str().unwrap(),
            )
            .unwrap();

        // Validate that everything matches
        assert!(
            dir_diff::is_different(testdata.path().join("testdst"), dst.path()).unwrap() == false
        );
    }
}