refactor: Make Output struct to help with indentation

src/gen/bash.rs

@@ -2,7 +2,13 @@ use std::{fs, path::Path};
 
 use anyhow::Result;
 
-use crate::{gen::Completions, parse::CommandInfo};
+use crate::{
+  gen::{
+    util::{quote_bash, Output},
+    Completions,
+  },
+  parse::CommandInfo,
+};
 
 pub struct BashCompletions;
 
@@ -15,16 +21,20 @@ impl Completions for BashCompletions {
     // TODO make option to not overwrite file
     let comp_name = format!("_comp_cmd_{cmd_name}");
 
-    let mut res = String::from("#!/usr/bin/env bash\n\n");
-    res.push_str(&format!("function {comp_name} {{\n"));
-    res.push_str("\tCOMPREPLY=()\n");
-    res.push_str("\tcase ${COMP_CWORD} in\n");
+    let mut res = Output::new(String::from("\t"));
+    res.writeln("#!/usr/bin/env bash\n");
+    res.writeln(&format!("function {} {{", comp_name));
+    res.writeln("COMPREPLY=()");
+    res.writeln("case ${COMP_CWORD} in");
     // generate_fn(&cmd_name, cmd_info, &mut res, 0, &comp_name);
-    res.push_str("\tesac\n");
-    res.push_str("\treturn 0\n");
-    res.push_str("}\n");
+    res.writeln("esac");
+    res.writeln("return 0");
+    res.writeln("}");
 
-    fs::write(out_dir.as_ref().join(format!("_{cmd_name}.bash")), res)?;
+    fs::write(
+      out_dir.as_ref().join(format!("_{}.bash", cmd_name)),
+      res.text(),
+    )?;
     Ok(())
   }
 }

src/gen/json.rs

@@ -2,9 +2,10 @@ use std::{fs, path::Path};
 
 use anyhow::Result;
 
-use crate::{gen::Completions, parse::CommandInfo};
-
-const INDENT: &str = "  ";
+use crate::{
+  gen::{util::Output, Completions},
+  parse::CommandInfo,
+};
 
 pub struct JsonCompletions;
 
@@ -16,11 +17,16 @@ impl Completions for JsonCompletions {
   where
     P: AsRef<Path>,
   {
-    let mut res = String::new();
-    res.push_str("{\n");
-    generate_cmd(&cmd_name, cmd_info, 1, true, &mut res);
-    res.push_str("}\n");
-    fs::write(out_dir.as_ref().join(format!("{cmd_name}.json")), res)?;
+    let mut res = Output::new(String::from("  "));
+    res.writeln("{");
+    res.indent();
+    generate_cmd(&cmd_name, cmd_info, true, &mut res);
+    res.dedent();
+    res.writeln("}");
+    fs::write(
+      out_dir.as_ref().join(format!("{}.json", cmd_name)),
+      res.text(),
+    )?;
     Ok(())
   }
 }
@@ -31,88 +37,77 @@ impl Completions for JsonCompletions {
 /// * `indent` - The indentation level (how many subcommands in we are)
 /// * `last` - Whether this is the last command at this level. Used for deciding
 ///   whether or not to put a trailing comma
-fn generate_cmd(
-  cmd: &str,
-  cmd_info: CommandInfo,
-  indent: usize,
-  last: bool,
-  out: &mut String,
-) {
+fn generate_cmd(cmd: &str, cmd_info: CommandInfo, last: bool, out: &mut Output) {
   let cmd = quote(cmd);
   // Avoid trailing commas
   let end = if last { "}" } else { "}," };
   let mut args = cmd_info.args.into_iter();
   if let Some(mut arg) = args.next() {
-    println_indent(indent, out, format!("{cmd}: {{"));
-    println_indent(indent + 1, out, "\"args\": [");
-    while {
-      println_indent(indent + 2, out, "{");
+    out.writeln(format!("{cmd}: {{"));
+    out.indent();
+    out.writeln("\"args\": [");
+    out.indent();
+
+    loop {
+      out.writeln("{");
+      out.indent();
+
       let forms = arg
         .forms
         .iter()
         .map(|a| quote(a))
         .collect::<Vec<_>>()
         .join(", ");
-      print_indent(indent + 3, out, format!(r#""forms": [{forms}]"#));
+      out.write(format!("\"forms\": [{}]", forms));
       if let Some(desc) = &arg.desc {
-        out.push_str(",\n");
-        println_indent(
-          indent + 3,
-          out,
-          format!(r#""description": {}"#, quote(desc)),
-        );
+        out.writeln(",");
+        out.writeln(format!("\"description\": {}", quote(desc)));
       } else {
-        out.push('\n');
+        out.writeln("");
       }
+
+      out.dedent();
       if let Some(next) = args.next() {
-        println_indent(indent + 2, out, "},");
+        out.writeln("},");
         arg = next;
-        true
       } else {
         // Avoid trailing comma
-        println_indent(indent + 2, out, "}");
-        false
+        out.writeln("}");
+        break;
       }
-    } {}
-    println_indent(indent + 1, out, "],");
+    }
+
+    out.dedent();
+    out.writeln("],");
 
     let mut subcmds = cmd_info.subcommands.into_iter();
     if let Some((mut name, mut info)) = subcmds.next() {
-      println_indent(indent + 1, out, "\"subcommands\": {");
+      out.writeln("\"subcommands\": {");
+      out.indent();
       loop {
         if let Some(next) = subcmds.next() {
-          generate_cmd(&name, info, indent + 2, false, out);
+          generate_cmd(&name, info, false, out);
           name = next.0;
           info = next.1;
         } else {
-          generate_cmd(&name, info, indent + 2, true, out);
+          generate_cmd(&name, info, true, out);
           break;
         }
       }
-      println_indent(indent + 1, out, "}");
+      out.dedent();
+      out.writeln("}");
     } else {
-      println_indent(indent + 1, out, "\"subcommands\": {}");
+      out.writeln("\"subcommands\": {}");
     }
 
-    println_indent(indent, out, end);
+    out.dedent();
+    out.writeln(end);
   } else {
     // If no arguments, print `"cmd": {}` on a single line
-    println_indent(indent, out, format!("{cmd}: {{{end}"))
+    out.writeln(format!("{}: {{{}", cmd, end));
   }
 }
 
 fn quote(s: &str) -> String {
   format!("\"{}\"", s.replace('\\', r"\\").replace('"', "\\\""))
 }
-
-/// Like print_indent, but with a newline
-fn println_indent<S: AsRef<str>>(indent: usize, out: &mut String, text: S) {
-  print_indent(indent, out, text);
-  out.push('\n');
-}
-
-/// Helper to print at a specific indentation level with a newline
-fn print_indent<S: AsRef<str>>(indent: usize, out: &mut String, text: S) {
-  out.push_str(&INDENT.repeat(indent));
-  out.push_str(text.as_ref());
-}

src/gen/mod.rs

@@ -1,6 +1,6 @@
 mod bash;
 mod json;
-pub(self) mod util;
+mod util;
 mod zsh;
 
 use std::path::Path;

src/gen/util.rs

@@ -1,5 +1,60 @@
 /// Wrap in single quotes (and escape single quotes inside) so that it's safe
 /// for Bash and Zsh to read
-pub fn quote(s: &str) -> String {
-  format!("'{}'", s.replace('\'', r#"'"'"'"#))
+pub fn quote_bash<S: AsRef<str>>(s: S) -> String {
+  format!("'{}'", s.as_ref().replace('\'', r#"'"'"'"#))
+}
+
+/// Helper to write indented text to a string
+pub struct Output {
+  text: String,
+  indent_str: String,
+  indent_level: usize,
+  /// If true, need to indent when the next string is written
+  line_ended: bool,
+}
+
+impl Output {
+  /// Make a new [Output]. `indent_str` is the string to indent with (e.g.
+  /// `"\t"`).
+  pub fn new(indent_str: String) -> Output {
+    Output {
+      text: String::from(""),
+      indent_str,
+      indent_level: 0,
+      line_ended: false,
+    }
+  }
+
+  /// Increase the indentation level by 1
+  pub fn indent(&mut self) {
+    self.indent_level += 1;
+  }
+
+  /// Decrease the indentation level by 1
+  pub fn dedent(&mut self) {
+    self.indent_level -= 1;
+  }
+
+  /// Write some text (without a newline)
+  pub fn write<S: AsRef<str>>(&mut self, s: S) {
+    if self.line_ended {
+      for _ in 0..self.indent_level {
+        self.text.push_str(&self.indent_str);
+      }
+      self.line_ended = false;
+    }
+    self.text.push_str(s.as_ref());
+  }
+
+  /// Write some text (with a newline)
+  pub fn writeln<S: AsRef<str>>(&mut self, s: S) {
+    self.write(s);
+    self.text.push('\n');
+    self.line_ended = true;
+  }
+
+  /// Consume this [Output], returning the text written to it
+  pub fn text(self) -> String {
+    self.text
+  }
 }

src/gen/zsh.rs

@@ -2,14 +2,12 @@ use std::{fs, path::Path};
 
 use anyhow::Result;
 
+use super::util::Output;
 use crate::{
   gen::{util, Completions},
   parse::CommandInfo,
 };
 
-/// Indentation to use (for readability)
-const INDENT: &str = "    ";
-
 pub struct ZshCompletions;
 
 impl Completions for ZshCompletions {
@@ -49,10 +47,14 @@ impl Completions for ZshCompletions {
     P: AsRef<Path>,
   {
     // TODO make option to not overwrite file
-    let comp_name = format!("_{cmd_name}");
-    let mut res = format!("#compdef {comp_name} {cmd_name}\n");
+    let comp_name = format!("_{}", cmd_name);
+    let mut res = Output::new(String::from("\t"));
+    res.writeln(format!("#compdef {} {}", comp_name, cmd_name));
     generate_fn(&cmd_name, cmd_info, &mut res, 0, &comp_name);
-    fs::write(out_dir.as_ref().join(format!("{comp_name}.zsh")), res)?;
+    fs::write(
+      out_dir.as_ref().join(format!("{}.zsh", comp_name)),
+      res.text(),
+    )?;
     Ok(())
   }
 }
@@ -68,26 +70,30 @@ impl Completions for ZshCompletions {
 fn generate_fn(
   _cmd_name: &str,
   cmd_info: CommandInfo,
-  out: &mut String,
+  out: &mut Output,
   pos: usize,
   fn_name: &str,
 ) {
-  out.push('\n');
-  out.push_str(&format!("function {fn_name} {{\n"));
+  out.write("\n");
+  out.writeln(format!("function {} {{", fn_name));
+  out.indent();
+
   if !cmd_info.subcommands.is_empty() {
-    out.push_str(&format!("{}{}", INDENT, "local line\n"));
+    out.writeln("local line");
   }
   if cmd_info.subcommands.is_empty() {
-    out.push_str(&format!("{INDENT}_arguments"));
+    out.write("_arguments");
   } else {
-    out.push_str(&format!("{INDENT}_arguments -C"));
+    out.write("_arguments -C");
   }
+
+  out.indent();
   for opt in cmd_info.args {
     let desc = opt.desc.unwrap_or_default();
     for form in opt.forms {
-      let text = util::quote(&format!("{form}[{}]", desc));
-      out.push_str(" \\\n");
-      out.push_str(&format!("{INDENT}{INDENT}{text}"));
+      let text = util::quote_bash(format!("{}[{}]", form, desc));
+      out.writeln(" \\");
+      out.write(text);
     }
   }
 
@@ -98,30 +104,32 @@ fn generate_fn(
       .map(|s| s.to_string())
       .collect::<Vec<_>>()
       .join(" ");
-    out.push_str(&format!(" \\\n{INDENT}{INDENT}': :({sub_cmds})'"));
-    out.push_str(&format!(" \\\n{INDENT}{INDENT}'*::arg:->args'\n"));
+    out.writeln(" \\");
+    out.writeln(format!("': :({})' \\", sub_cmds));
+    out.writeln("'*::arg:->args'");
+    out.dedent();
 
-    out.push_str(&format!("{INDENT}case $line[{}] in\n", pos + 1));
+    out.writeln(format!("case $line[{}] in", pos + 1));
+    out.indent();
     for sub_cmd in cmd_info.subcommands.keys() {
-      out.push_str(&format!(
-        "{INDENT}{INDENT}{sub_cmd}) {}_{};;\n",
-        fn_name, sub_cmd
-      ))
+      out.writeln(format!("{sub_cmd}) {}_{};;", fn_name, sub_cmd))
     }
-    out.push_str(&format!("{INDENT}esac\n"));
+    out.dedent();
+    out.writeln("esac");
   } else {
-    out.push('\n');
+    out.write("\n");
   }
 
-  out.push_str("}\n");
+  out.dedent();
+  out.writeln("}");
 
   for (sub_cmd, sub_cmd_info) in cmd_info.subcommands {
     generate_fn(
       &sub_cmd,
       sub_cmd_info,
       out,
       pos + 1,
-      &format!("{fn_name}_{sub_cmd}"),
+      &format!("{}_{}", fn_name, sub_cmd),
     );
   }
 }