From ed181a65180ca758073623874dd352b15df6f027 Mon Sep 17 00:00:00 2001 From: Joscha Date: Wed, 10 Aug 2022 01:58:25 +0200 Subject: [PATCH] Restructure export code and arg handling --- src/export.rs | 161 +++++++++++++++++++++++++-------------------- src/export/text.rs | 94 ++++++++++++++++++++++++++ src/main.rs | 4 +- 3 files changed, 186 insertions(+), 73 deletions(-) create mode 100644 src/export/text.rs diff --git a/src/export.rs b/src/export.rs index 52c6934..2430865 100644 --- a/src/export.rs +++ b/src/export.rs @@ -1,96 +1,115 @@ //! Export logs from the vault to plain text files. +mod text; + use std::fs::File; use std::io::{BufWriter, Write}; -use std::path::Path; -use time::format_description::FormatItem; -use time::macros::format_description; -use unicode_width::UnicodeWidthStr; - -use crate::euph::api::Snowflake; -use crate::euph::SmallMessage; -use crate::store::{MsgStore, Tree}; use crate::vault::Vault; -const TIME_FORMAT: &[FormatItem<'_>] = - format_description!("[year]-[month]-[day] [hour]:[minute]:[second]"); -const TIME_EMPTY: &str = " "; - -pub async fn export(vault: &Vault, room: String, file: &Path) -> anyhow::Result<()> { - println!("Exporting &{room} to {}", file.to_string_lossy()); - let mut file = BufWriter::new(File::create(file)?); - let vault = vault.euph(room); - - let mut exported_trees = 0; - let mut exported_msgs = 0; - let mut tree_id = vault.first_tree_id().await; - while let Some(some_tree_id) = tree_id { - let tree = vault.tree(&some_tree_id).await; - write_tree(&mut file, &tree, some_tree_id, 0)?; - tree_id = vault.next_tree_id(&some_tree_id).await; - - exported_trees += 1; - exported_msgs += tree.len(); - - if exported_trees % 10000 == 0 { - println!("Exported {exported_trees} trees, {exported_msgs} messages") - } - } - println!("Exported {exported_trees} trees, {exported_msgs} messages in total"); - - file.flush()?; - Ok(()) +#[derive(Debug, Clone, Copy, clap::ValueEnum)] +pub enum Format { + /// Human-readable tree-structured messages. + Text, } -fn write_tree( - file: &mut BufWriter, - tree: &Tree, - id: Snowflake, - indent: usize, -) -> anyhow::Result<()> { - let indent_string = "| ".repeat(indent); +impl Format { + fn name(&self) -> &'static str { + match self { + Self::Text => "text", + } + } - if let Some(msg) = tree.msg(&id) { - write_msg(file, &indent_string, msg)?; + fn extension(&self) -> &'static str { + match self { + Self::Text => "txt", + } + } +} + +#[derive(Debug, clap::Parser)] +pub struct Args { + rooms: Vec, + + /// Export all rooms. + #[clap(long, short)] + all: bool, + + /// Format of the output file. + #[clap(long, short, value_enum, default_value_t = Format::Text)] + format: Format, + + /// Location of the output file + /// + /// May include the following placeholders: + /// `%r` - room name + /// `%e` - format extension + /// A literal `%` can be written as `%%`. + /// + /// If the value ends with a `/`, it is assumed to point to a directory and + /// `%r.%e` will be appended. + /// + /// Must be a valid utf-8 encoded string. + #[clap(long, short, default_value_t = Into::into("%r.%e"))] + #[clap(verbatim_doc_comment)] + out: String, +} + +pub async fn export(vault: &Vault, mut args: Args) -> anyhow::Result<()> { + if args.out.ends_with('/') { + args.out.push_str("%r.%e"); + } + + let rooms = if args.all { + let mut rooms = vault.euph_rooms().await; + rooms.sort_unstable(); + rooms } else { - write_placeholder(file, &indent_string)?; + let mut rooms = args.rooms.clone(); + rooms.dedup(); + rooms + }; + + if rooms.is_empty() { + println!("No rooms to export"); } - if let Some(children) = tree.children(&id) { - for child in children { - write_tree(file, tree, *child, indent + 1)?; + for room in rooms { + let out = format_out(&args.out, &room, args.format); + println!("Exporting &{room} as {} to {out}", args.format.name()); + + let mut file = BufWriter::new(File::create(out)?); + match args.format { + Format::Text => text::export_to_file(vault, room, &mut file).await?, } + file.flush()?; } Ok(()) } -fn write_msg( - file: &mut BufWriter, - indent_string: &str, - msg: &SmallMessage, -) -> anyhow::Result<()> { - let nick = &msg.nick; - let nick_empty = " ".repeat(nick.width()); +fn format_out(out: &str, room: &str, format: Format) -> String { + let mut result = String::new(); - for (i, line) in msg.content.lines().enumerate() { - if i == 0 { - let time = msg - .time - .0 - .format(TIME_FORMAT) - .expect("time can be formatted"); - writeln!(file, "{time} {indent_string}[{nick}] {line}")?; + let mut special = false; + for char in out.chars() { + if special { + match char { + 'r' => result.push_str(room), + 'e' => result.push_str(format.extension()), + '%' => result.push('%'), + _ => { + result.push('%'); + result.push(char); + } + } + special = false; + } else if char == '%' { + special = true; } else { - writeln!(file, "{TIME_EMPTY} {indent_string}| {nick_empty} {line}")?; + result.push(char); } } - Ok(()) -} - -fn write_placeholder(file: &mut BufWriter, indent_string: &str) -> anyhow::Result<()> { - writeln!(file, "{TIME_EMPTY} {indent_string}[...]")?; - Ok(()) + result } diff --git a/src/export/text.rs b/src/export/text.rs new file mode 100644 index 0000000..7af4e89 --- /dev/null +++ b/src/export/text.rs @@ -0,0 +1,94 @@ +use std::fs::File; +use std::io::{BufWriter, Write}; + +use time::format_description::FormatItem; +use time::macros::format_description; +use unicode_width::UnicodeWidthStr; + +use crate::euph::api::Snowflake; +use crate::euph::SmallMessage; +use crate::store::{MsgStore, Tree}; +use crate::vault::Vault; + +const TIME_FORMAT: &[FormatItem<'_>] = + format_description!("[year]-[month]-[day] [hour]:[minute]:[second]"); +const TIME_EMPTY: &str = " "; + +pub async fn export_to_file( + vault: &Vault, + room: String, + file: &mut BufWriter, +) -> anyhow::Result<()> { + let vault = vault.euph(room); + + let mut exported_trees = 0; + let mut exported_msgs = 0; + let mut tree_id = vault.first_tree_id().await; + while let Some(some_tree_id) = tree_id { + let tree = vault.tree(&some_tree_id).await; + write_tree(file, &tree, some_tree_id, 0)?; + tree_id = vault.next_tree_id(&some_tree_id).await; + + exported_trees += 1; + exported_msgs += tree.len(); + + if exported_trees % 10000 == 0 { + println!(" {exported_trees} trees, {exported_msgs} messages") + } + } + println!(" {exported_trees} trees, {exported_msgs} messages in total"); + + Ok(()) +} + +fn write_tree( + file: &mut BufWriter, + tree: &Tree, + id: Snowflake, + indent: usize, +) -> anyhow::Result<()> { + let indent_string = "| ".repeat(indent); + + if let Some(msg) = tree.msg(&id) { + write_msg(file, &indent_string, msg)?; + } else { + write_placeholder(file, &indent_string)?; + } + + if let Some(children) = tree.children(&id) { + for child in children { + write_tree(file, tree, *child, indent + 1)?; + } + } + + Ok(()) +} + +fn write_msg( + file: &mut BufWriter, + indent_string: &str, + msg: &SmallMessage, +) -> anyhow::Result<()> { + let nick = &msg.nick; + let nick_empty = " ".repeat(nick.width()); + + for (i, line) in msg.content.lines().enumerate() { + if i == 0 { + let time = msg + .time + .0 + .format(TIME_FORMAT) + .expect("time can be formatted"); + writeln!(file, "{time} {indent_string}[{nick}] {line}")?; + } else { + writeln!(file, "{TIME_EMPTY} {indent_string}| {nick_empty} {line}")?; + } + } + + Ok(()) +} + +fn write_placeholder(file: &mut BufWriter, indent_string: &str) -> anyhow::Result<()> { + writeln!(file, "{TIME_EMPTY} {indent_string}[...]")?; + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index e940e5e..31d72d2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -43,7 +43,7 @@ enum Command { /// Run the client interactively (default). Run, /// Export logs for a single room as a plain text file. - Export { room: String, file: PathBuf }, + Export(export::Args), /// Compact and clean up vault. Gc, /// Clear euphoria session cookies. @@ -87,7 +87,7 @@ async fn main() -> anyhow::Result<()> { match args.command.unwrap_or_default() { Command::Run => run(&vault, args.measure_widths).await?, - Command::Export { room, file } => export::export(&vault, room, &file).await?, + Command::Export(args) => export::export(&vault, args).await?, Command::Gc => { println!("Cleaning up and compacting vault"); println!("This may take a while...");