Support prefix

README.md

@@ -70,6 +70,12 @@ fn gcd_list(list: &[u64]) -> u64 {
     list.iter().fold(list[0], |a, &b| gcd(a, b))
 }
 
+// You can set prefix string.
+// Note: All codes will be formatted by rustfmt on output
+#[snippet(prefix = "use std::io::{self,Read};")]
+#[snippet(prefix = "use std::str::FromStr;")]
+fn foo() {}
+
 #[test]
 fn test_gcd() {
     assert_eq!(gcd(57, 3), 3);
@@ -91,6 +97,11 @@ Extract snippet !
 
 ```
 $ cargo snippet
+snippet foo
+    use std::io::{self, Read};
+    use std::str::FromStr;
+    fn foo() {}
+
 snippet gcd
     fn gcd(a: u64, b: u64) -> u64 {
         if b == 0 {
@@ -109,7 +120,7 @@ snippet gcd_list
         }
     }
     fn gcd_list(list: &[u64]) -> u64 {
-        list.iter().fold(list[0], |a, b| gcd(a, b));
+        list.iter().fold(list[0], |a, &b| gcd(a, b))
     }
 
 snippet lcm

src/parser.rs

@@ -198,6 +198,43 @@ fn get_snippet_uses(attr: &Attribute) -> Option<Vec<String>> {
     })
 }
 
+fn get_simple_attr(attr: &Attribute, key: &str) -> Vec<String> {
+    attr.parse_meta()
+        .ok()
+        .and_then(|metaitem| {
+            if !is_snippet_path(metaitem.path().to_token_stream().to_string().as_str()) {
+                return None;
+            }
+
+            match metaitem {
+                // #[snippet(`key`="..")]
+                Meta::List(list) => list
+                    .nested
+                    .iter()
+                    .filter_map(|item| {
+                        if let NestedMeta::Meta(Meta::NameValue(ref nv)) = item {
+                            if nv.path.to_token_stream().to_string() == key {
+                                let value = if let syn::Lit::Str(s) = &nv.lit.clone() {
+                                    s.value()
+                                } else {
+                                    panic!("attribute must be string");
+                                };
+                                Some(value)
+                            } else {
+                                None
+                            }
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<_>>()
+                    .into(),
+                _ => None,
+            }
+        })
+        .unwrap_or(Vec::new())
+}
+
 fn parse_attrs(
     attrs: &[Attribute],
     default_snippet_name: Option<String>,
@@ -246,7 +283,18 @@ fn parse_attrs(
         .flat_map(|v| v.into_iter())
         .collect::<HashSet<_>>();
 
-    Some(SnippetAttributes { names, uses })
+    let prefix = attrs
+        .iter()
+        .map(|attr| get_simple_attr(attr, "prefix").into_iter())
+        .flatten()
+        .collect::<Vec<_>>()
+        .join("\n");
+
+    Some(SnippetAttributes {
+        names,
+        uses,
+        prefix,
+    })
 }
 
 // Get snippet names and snippet code (not formatted)
@@ -631,10 +679,72 @@ mod test {
         "#;
 
         let snip = snippets(&src);
-        dbg!(&snip);
         assert_eq!(
             format_src(snip["bar"].as_str()).unwrap(),
             format_src("fn bar() {}").unwrap()
         );
     }
+
+    #[test]
+    fn test_attribute_prefix() {
+        let src = r#"
+            #[snippet(prefix = "use std::io;")]
+            fn bar() {}
+        "#;
+
+        let snip = snippets(&src);
+        assert_eq!(
+            format_src(snip["bar"].as_str()).unwrap(),
+            format_src("use std::io;\nfn bar() {}").unwrap()
+        );
+
+        let src = r#"
+            #[snippet(prefix="use std::io::{self,Read};\nuse std::str::FromStr;")]
+            fn bar() {}
+        "#;
+
+        let snip = snippets(&src);
+        assert_eq!(
+            format_src(snip["bar"].as_str()).unwrap(),
+            format_src("use std::io::{self,Read};\nuse std::str::FromStr;\nfn bar() {}").unwrap()
+        );
+
+        let src = r#"
+            #[snippet(prefix=r"use std::io::{self,Read};
+use std::str::FromStr;")]
+            fn bar() {}
+        "#;
+
+        let snip = snippets(&src);
+        dbg!(&snip);
+        assert_eq!(
+            format_src(snip["bar"].as_str()).unwrap(),
+            format_src("use std::io::{self,Read};\nuse std::str::FromStr;\nfn bar() {}").unwrap()
+        );
+    }
+
+    #[test]
+    fn test_attribute_prefix_include() {
+        let src = r#"
+            #[snippet(prefix = "use std::sync;")]
+            fn foo() {}
+            #[snippet(prefix = "use std::io;", include = "foo")]
+            fn bar() {}
+        "#;
+
+        let snip = snippets(&src);
+        assert_eq!(
+            format_src(snip["bar"].as_str()).unwrap(),
+            format_src(
+                &quote!(
+                    use std::sync;
+                    use std::io;
+                    fn foo() {}
+                    fn bar() {}
+                )
+                .to_string()
+            )
+            .unwrap()
+        );
+    }
 }

src/snippet.rs

@@ -5,6 +5,8 @@ pub struct SnippetAttributes {
     pub names: HashSet<String>,
     // Dependencies
     pub uses: HashSet<String>,
+    // Prefix for snippet. It's will be emitted prior to the snippet.
+    pub prefix: String,
 }
 
 pub struct Snippet {
@@ -14,12 +16,20 @@ pub struct Snippet {
 }
 
 pub fn process_snippets(snips: &[Snippet]) -> BTreeMap<String, String> {
-    let mut pre: BTreeMap<String, String> = BTreeMap::new();
+    #[derive(Default, Clone, Debug)]
+    struct Snip {
+        prefix: String,
+        content: String,
+    }
+
+    let mut pre: BTreeMap<String, Snip> = BTreeMap::new();
     let mut deps: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
 
     for snip in snips {
         for name in &snip.attrs.names {
-            *pre.entry(name.clone()).or_insert_with(String::new) += &snip.content;
+            let s = pre.entry(name.clone()).or_default();
+            s.prefix += &snip.attrs.prefix;
+            s.content += &snip.content;
 
             for dep in &snip.attrs.uses {
                 deps.entry(name.clone())
@@ -29,7 +39,7 @@ pub fn process_snippets(snips: &[Snippet]) -> BTreeMap<String, String> {
         }
     }
 
-    let mut res: BTreeMap<String, String> = BTreeMap::new();
+    let mut res: BTreeMap<String, Snip> = BTreeMap::new();
 
     for (name, uses) in &deps {
         let mut used = HashSet::new();
@@ -40,7 +50,10 @@ pub fn process_snippets(snips: &[Snippet]) -> BTreeMap<String, String> {
             if !used.contains(&dep) {
                 used.insert(dep.clone());
                 if let Some(c) = &pre.get(&dep) {
-                    *res.entry(name.clone()).or_insert_with(String::new) += c.as_str();
+                    // *res.entry(name.clone()).or_insert_with(String::new) += c.as_str();
+                    let s = res.entry(name.clone()).or_default();
+                    s.prefix += &c.prefix;
+                    s.content += &c.content;
 
                     if let Some(ds) = deps.get(&dep) {
                         for d in ds {
@@ -56,10 +69,14 @@ pub fn process_snippets(snips: &[Snippet]) -> BTreeMap<String, String> {
         }
     }
 
-    for (name, content) in pre {
+    for (name, snip) in pre {
         // Dependency first
-        *res.entry(name).or_insert_with(String::new) += content.as_str();
+        let s = res.entry(name).or_default();
+        s.prefix += snip.prefix.as_str();
+        s.content += snip.content.as_str();
     }
 
-    res
+    res.into_iter()
+        .map(|(k, v)| (k, v.prefix + v.content.as_str()))
+        .collect()
 }