feat: support adding newlines to files

Signed-off-by: Daniel Silverstone <dsilvers@digital-scurf.org>

1 files changed, 25 insertions, 7 deletions

diff --git a/src/ast.rs b/src/ast.rs
index 012e05a..f627178 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -970,11 +970,8 @@ pub struct DataFile {
 }
 
 impl DataFile {
-    fn new(filename: &str, contents: &str) -> DataFile {
-        DataFile {
-            filename: filename.to_string(),
-            contents: contents.to_string(),
-        }
+    fn new(filename: String, contents: String) -> DataFile {
+        DataFile { filename, contents }
     }
 
     pub fn filename(&self) -> &str {
@@ -1010,8 +1007,18 @@ impl MutVisitor for DataFiles {
             match block {
                 Block::CodeBlock(attr, contents) => {
                     if is_class(attr, "file") {
-                        self.files
-                            .push(DataFile::new(&get_filename(attr), &contents));
+                        let add_newline = match find_attr_kv(&attr, "add-newline").next() {
+                            None | Some("auto") => !contents.ends_with('\n'),
+                            Some("yes") => true,
+                            Some("no") => false,
+                            _ => unreachable!(),
+                        };
+                        let contents = if add_newline {
+                            format!("{}\n", contents)
+                        } else {
+                            contents.clone()
+                        };
+                        self.files.push(DataFile::new(get_filename(attr), contents));
                     }
                 }
                 _ => {
@@ -1081,3 +1088,14 @@ pub fn get_basedir_from(filename: &Path) -> Result<PathBuf> {
     };
     Ok(dirname)
 }
+
+/// Utility function to find key/value pairs from an attribute
+fn find_attr_kv<'a>(attr: &'a Attr, key: &'static str) -> impl Iterator<Item = &'a str> {
+    attr.2.iter().flat_map(move |(key_, value)| {
+        if key == key_ {
+            Some(value.as_ref())
+        } else {
+            None
+        }
+    })
+}