From cdf28edec6c801334c0df658e5ef321090aae3f4 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Sun, 12 Dec 2021 20:35:57 +0000
Subject: [PATCH] Add --reformat and --reformat-all flags

---
 src/cli.rs   | 13 ++++++++++++-
 src/files.rs |  6 ++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/cli.rs b/src/cli.rs
index 62dc7ce..3e714eb 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -18,6 +18,12 @@ pub struct Opt {
     /// File to load
     #[structopt(short, long, parse(from_os_str))]
     file: Option<PathBuf>,
+    /// Reformat the file
+    #[structopt(short, long)]
+    reformat: bool,
+    /// Reformat the file and all imports
+    #[structopt(short = "R", long)]
+    reformat_all: bool,
 }
 
 fn default_file() -> PathBuf {
@@ -31,7 +37,7 @@ pub fn run() -> anyhow::Result<()> {
     let opt = Opt::from_args();
 
     let file = opt.file.unwrap_or_else(default_file);
-    let files = Files::load(&file)?;
+    let mut files = Files::load(&file)?;
     let now = files.now().naive_local();
 
     let range = DateRange::new(
@@ -49,6 +55,11 @@ pub fn run() -> anyhow::Result<()> {
     render.render(&files, &entries, &layout);
     print!("{}", render.display());
 
+    if opt.reformat_all {
+        files.mark_all_dirty();
+    } else if opt.reformat {
+        files.mark_main_dirty();
+    }
     files.save()?;
     Ok(())
 }
diff --git a/src/files.rs b/src/files.rs
index bea36a3..86c3443 100644
--- a/src/files.rs
+++ b/src/files.rs
@@ -175,6 +175,12 @@ impl Files {
         Ok(())
     }
 
+    pub fn mark_main_dirty(&mut self) {
+        if let Some(file) = self.files.get_mut(0) {
+            file.dirty = true;
+        }
+    }
+
     pub fn mark_all_dirty(&mut self) {
         for file in self.files.iter_mut() {
             file.dirty = true;