diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/main.rs | 340 |
1 files changed, 340 insertions, 0 deletions
diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..2fe5a26 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,340 @@ +use std::collections::HashMap; +use std::ffi::{OsStr, OsString}; +use std::fmt; +use std::fs; +use std::io; +use std::path::PathBuf; +use std::process; + +use rand::{thread_rng, seq::{SliceRandom, IteratorRandom}}; +use clap::Parser; + +#[derive(Debug)] +enum Error { + NoSentences, + MultipleSentences, + BadTemplate(String), + IOFail(io::Error), + OutOfSentences, + OutOfWords(OsString), + UnknownCategory(OsString), + NoCategories, + DuplicateCategory(OsString), + BadPath(PathBuf), + EmptyOption, +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Error::NoSentences => + write!(f, "no sentences.txt file found"), + Error::MultipleSentences => + write!(f, "multiple sentences.txt files found, i'll be real i don't know how this is possible"), + Error::BadTemplate(t) => + write!(f, "unable to parse template sentence: {}", t), + Error::IOFail(e) => + write!(f, "IO error: {:?}", e), + Error::OutOfSentences => + write!(f, "ran out of sentences"), + Error::OutOfWords(c) => + write!(f, "ran out of words in category: {:?}", c), + Error::UnknownCategory(c) => + write!(f, "tried to pick word in unknown category: {:?}", c), + Error::NoCategories => + write!(f, "no categories provided"), + Error::DuplicateCategory(c) => + write!(f, "tried to create a category twice: {:?}", c), + Error::BadPath(p) => + write!(f, "failed to parse path: {:?}", p), + Error::EmptyOption => + write!(f, "multi-option variable had an empty option"), + } + } +} + +impl From<io::Error> for Error { + fn from(error: io::Error) -> Self { + Error::IOFail(error) + } +} + +type ProgResult<T> = Result<T, Error>; + +#[derive(Parser, Debug)] +#[command(name = "spout")] +#[command(about = "Generates nonsense", long_about = None)] +struct Args { + /// Path to data directory + directory: String, + + /// Number of sentences to generate + #[arg(short = 'n', long = "sentences", default_value_t = 1)] + n_sentences: u32, + + /// Don't re-use entries in data files (program will fail upon running out) + #[arg(long)] + unique: bool, +} + +#[derive(Debug, Clone)] +enum DropIn { + Basic(String), + Var(String), + OneOf(Vec<DropIn>), +} + +impl DropIn { + pub fn parse(text: &str) -> ProgResult<Self> { + if text.is_empty() { + return Err(Error::EmptyOption) + } + + if text.contains('|') { + let mut option_strs = text.split('|').collect::<Vec<&str>>(); + option_strs.dedup(); + let options = option_strs.into_iter() + .map(DropIn::parse) + .collect::<Result<Vec<_>, _>>()?; + return Ok(DropIn::OneOf(options)) + } + + if let Some(var_name) = text.strip_prefix('?') { + return Ok(DropIn::Var(var_name.to_string())); + } + + Ok(DropIn::Basic(text.to_string())) + } +} + +#[derive(Debug)] +struct CategorySet { + var_table: HashMap<String, OsString>, + categories: HashMap<OsString, Vec<String>>, + unique: bool, +} + +impl CategorySet { + pub fn new(unique: bool) -> Self { + CategorySet { + var_table: HashMap::new(), + categories: HashMap::new(), + unique, + } + } + + fn resolve_variable(&mut self, var: &str) -> ProgResult<OsString> { + if !self.var_table.contains_key(var) { + let new_cat = self.random_category()?; + self.var_table.insert(var.to_string(), new_cat); + } + Ok(self.var_table.get(var).unwrap().clone()) + } + + pub fn add_category(&mut self, name: OsString, mut words: Vec<String>) -> ProgResult<()> { + words.retain(|s| !s.is_empty()); + words.shuffle(&mut thread_rng()); + match self.categories.insert(name.clone(), words) { + None => Ok(()), + Some(_) => Err(Error::DuplicateCategory(name)), + } + } + + // not affected by --unique + pub fn random_category(&self) -> ProgResult<OsString> { + self.categories.keys() + .choose(&mut thread_rng()) + .ok_or(Error::NoCategories) + .map(|c| c.clone()) + } + + pub fn random_from_drop_in(&mut self, drop_in: &DropIn) -> ProgResult<String> { + let simple_di = if let DropIn::OneOf(ds) = drop_in { + ds.choose(&mut thread_rng()) + .expect("got a OneOf that is empty, somehow") + } else { + &drop_in + }; + + let category = match simple_di { + DropIn::Basic(s) => { + OsString::from(s) + }, + DropIn::Var(s) => { + self.resolve_variable(&s)? + }, + DropIn::OneOf(_) => { + panic!("Nested OneOf constructs are not supported") + }, + }; + + let words = self.categories.get_mut(&category) + .ok_or(Error::UnknownCategory(category.clone()))?; + + if self.unique { + words.pop() + .ok_or(Error::OutOfWords(category)) + } else { + words.choose(&mut thread_rng()) + .ok_or(Error::OutOfWords(category)) + .map(|c| c.clone()) + } + } +} + +#[derive(Debug, Clone)] +struct Sentence { + text_fragments: Vec<String>, + drop_ins: Vec<DropIn>, +} + +impl Sentence { + pub fn parse(original: String) -> ProgResult<Self> { + let mut text = original.as_str(); + let mut text_fragments = vec![]; + let mut drop_ins = vec![]; + + while let Some((front, back)) = text.split_once("%%") { + let (drop_in_body, rest) = back.split_once("%%") + .ok_or(Error::BadTemplate(original.clone()))?; + let drop_in = DropIn::parse(drop_in_body)?; + + text_fragments.push(front.to_string()); + drop_ins.push(drop_in); + text = rest; + } + text_fragments.push(text.to_string()); + + Ok(Sentence { + text_fragments, + drop_ins, + }) + } + + pub fn generate(&self, categories: &mut CategorySet) -> ProgResult<String> { + let mut text_iter = self.text_fragments.iter(); + let mut drop_ins_iter = self.drop_ins.iter(); + let mut buffer = String::new(); + + while let Some(drop_in) = drop_ins_iter.next() { + let text = text_iter.next() + .expect("invariant: number of text = number of drop-in + 1"); + buffer.push_str(text); + + let drop_in_text = categories.random_from_drop_in(&drop_in)?; + buffer.push_str(drop_in_text.as_str()); + } + + let end_text = text_iter.next() + .expect("invariant: number of text = number of drop-in + 1"); + buffer.push_str(end_text); + + Ok(buffer) + } +} + +#[derive(Debug)] +struct SentenceSet { + sentences: Vec<Sentence>, + unique: bool, +} + +impl SentenceSet { + pub fn new(mut sentence_strings: Vec<String>, unique: bool) -> ProgResult<Self> { + sentence_strings.retain(|s| !s.is_empty()); + sentence_strings.shuffle(&mut thread_rng()); + + let sentences = sentence_strings.into_iter() + .map(Sentence::parse) + .collect::<Result<Vec<_>, _>>()?; + + Ok(SentenceSet { + sentences, + unique, + }) + } + + pub fn random_sentence(&mut self) -> ProgResult<Sentence> { + if self.unique { + self.sentences.pop() + .ok_or(Error::OutOfSentences) + } else { + self.sentences.choose(&mut thread_rng()) + .ok_or(Error::OutOfSentences) + .map(|s| s.clone()) + } + } +} + +fn read_lines(path: &PathBuf) -> ProgResult<Vec<String>> { + let mut lines = vec![]; + + let text = fs::read_to_string(path)?; + for line in text.lines() { + lines.push(line.to_string()); + } + + Ok(lines) +} + +fn read_files(dir_path: &str, unique: bool) -> ProgResult<(SentenceSet, CategorySet)> { + let mut opt_sentences: Option<SentenceSet> = None; + let mut categories = CategorySet::new(unique); + + let dir = fs::read_dir(dir_path)?; + for entry in dir { + let path_buf = entry?.path(); + if path_buf.is_file() { + let extension = path_buf.extension() + .ok_or(Error::BadPath(path_buf.clone()))?; + if extension == OsStr::new("txt") { + let lines = read_lines(&path_buf)?; + let filename = path_buf.file_name() + .ok_or(Error::BadPath(path_buf.clone()))?; + if filename == OsStr::new("sentences.txt") { + if opt_sentences.is_none() { + opt_sentences = Some(SentenceSet::new(lines, unique)?); + } else { + return Err(Error::MultipleSentences); + } + } else { + let cat_name = path_buf.file_stem() + .expect("filename stem unwrapping osstr") + .to_os_string(); + categories.add_category(cat_name, lines)?; + } + } + } + } + + match opt_sentences { + None => Err(Error::NoSentences), + Some(s) => Ok((s, categories)), + } +} + +fn crash(e: Error) -> ! { + eprintln!("[ERROR] {}", e); + process::exit(1); +} + +fn sub_main(args: Args) -> ProgResult<()> { + let (mut sentences, mut categories) = read_files(&args.directory, args.unique)?; + + for _i in 0..args.n_sentences { + let sentence = sentences.random_sentence()?; + let generated = sentence.generate(&mut categories)?; + println!("{}", generated); + } + + Ok(()) +} + +fn main() { + let args = Args::parse(); + + if let Err(e) = sub_main(args) { + crash(e) + } +} + |