Merge pull request #34 from robertpyke/master

Add support for shebang

src/licensure.rs

@@ -1,6 +1,7 @@
 use std::fs::File;
-use std::io;
-use std::io::prelude::*;
+use std::io::{self, prelude::*};
+
+use regex::Regex;
 
 use crate::config::Config;
 
@@ -57,6 +58,26 @@ impl Licensure {
                 }
             }
 
+            let shebang_re: Regex = Regex::new(r"^#!.*\n").expect("shebang regex didn't compile!");
+
+            // Check for a shebang (scoped for content borrow)
+            let shebang_end = {
+                let shebang_match_opt = shebang_re.find(&content);
+                match shebang_match_opt {
+                    Some(shebang_match) => Option::Some(shebang_match.end()),
+                    None => Option::None,
+                }
+            };
+            // If we idenfied a shebang, strip it from content (we'll add it back at the end)
+            let shebang = match shebang_end {
+                Some(split_at) => {
+                    let mut shebang = content;
+                    content = shebang.split_off(split_at);
+                    Some(shebang)
+                }
+                None => Option::None,
+            };
+
             if content.contains(&header) {
                 info!("{} already licensed", file);
                 continue;
@@ -114,6 +135,14 @@ impl Licensure {
                 header.push_str(&content);
             }
 
+            // Put the shebang back (if we had one)
+            match shebang {
+                Some(val) => {
+                    header.insert_str(0, &val);
+                }
+                None => {}
+            }
+
             if self.config.change_in_place {
                 let mut f = File::create(file)?;
                 f.write_all(header.as_bytes())?;