Restructure export code and arg handling

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<'_>] =
+#[derive(Debug, Clone, Copy, clap::ValueEnum)]
-    format_description!("[year]-[month]-[day] [hour]:[minute]:[second]");
+pub enum Format {
-const TIME_EMPTY: &str = "                   ";
+    /// Human-readable tree-structured messages.
-
+    Text,
-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(())
 }
 
-fn write_tree(
+impl Format {
-    file: &mut BufWriter<File>,
+    fn name(&self) -> &'static str {
-    tree: &Tree<SmallMessage>,
+        match self {
-    id: Snowflake,
+            Self::Text => "text",
-    indent: usize,
+        }
-) -> anyhow::Result<()> {
+    }
-    let indent_string = "| ".repeat(indent);
 
-    if let Some(msg) = tree.msg(&id) {
+    fn extension(&self) -> &'static str {
-        write_msg(file, &indent_string, msg)?;
+        match self {
+            Self::Text => "txt",
+        }
+    }
+}
+
+#[derive(Debug, clap::Parser)]
+pub struct Args {
+    rooms: Vec<String>,
+
+    /// 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 room in rooms {
-        for child in children {
+        let out = format_out(&args.out, &room, args.format);
-            write_tree(file, tree, *child, indent + 1)?;
+        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(
+fn format_out(out: &str, room: &str, format: Format) -> String {
-    file: &mut BufWriter<File>,
+    let mut result = String::new();
-    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() {
+    let mut special = false;
-        if i == 0 {
+    for char in out.chars() {
-            let time = msg
+        if special {
-                .time
+            match char {
-                .0
+                'r' => result.push_str(room),
-                .format(TIME_FORMAT)
+                'e' => result.push_str(format.extension()),
-                .expect("time can be formatted");
+                '%' => result.push('%'),
-            writeln!(file, "{time} {indent_string}[{nick}] {line}")?;
+                _ => {
+                    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(())
+    result
-}
-
-fn write_placeholder(file: &mut BufWriter<File>, indent_string: &str) -> anyhow::Result<()> {
-    writeln!(file, "{TIME_EMPTY} {indent_string}[...]")?;
-    Ok(())
 }

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<File>,
+) -> 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<File>,
+    tree: &Tree<SmallMessage>,
+    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<File>,
+    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<File>, indent_string: &str) -> anyhow::Result<()> {
+    writeln!(file, "{TIME_EMPTY} {indent_string}[...]")?;
+    Ok(())
+}

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...");