proc_include_dir_as_map/
lib.rs

1#![forbid(unsafe_code)]
2
3use std::{
4    collections::HashMap,
5    fs::File,
6    io::Read,
7    path::{Path, PathBuf},
8};
9
10use proc_macro::TokenStream;
11use quote::quote;
12
13#[allow(dead_code)]
14type DirMap = HashMap<String, Vec<u8>>;
15
16#[allow(dead_code)]
17fn file_to_bytes(path: &Path) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
18    let mut data = Vec::new();
19    let mut file = File::open(path)?;
20    file.read_to_end(&mut data)?;
21    Ok(data)
22}
23
24#[allow(dead_code)]
25fn dir_to_map(root: &Path, base: &Path) -> Result<DirMap, Box<dyn std::error::Error>> {
26    let mut paths = HashMap::new();
27    for entry in std::fs::read_dir(base)? {
28        let entry = entry?;
29        let path = entry.path();
30        let metadata = std::fs::metadata(&path)?;
31        if metadata.is_file() {
32            let data = file_to_bytes(&path)?;
33            let rel = path.strip_prefix(root)?.to_str().unwrap();
34            paths.insert(rel.to_string(), data);
35        } else if metadata.is_dir() {
36            let dirmap = dir_to_map(root, &path)?;
37            paths.extend(dirmap);
38        } else {
39            panic!("{:?} is not a file or directory", entry);
40        }
41    }
42    Ok(paths)
43}
44
45fn env_expand_dir(raw: &str) -> PathBuf {
46    let mut copy = raw;
47    let mut root = String::new();
48    while let Some(pos) = copy.find('$') {
49        let (head, tail) = copy.split_at(pos);
50        let token = &tail[1..];
51        let end = token
52            .find(|ch: char| !ch.is_alphanumeric() && ch != '_')
53            .unwrap_or(token.len());
54        let env = &token[..end];
55        let val = std::env::var(env)
56            .unwrap_or_else(|_| panic!("{:?} is not a valid environment variable", env));
57        copy = token.strip_prefix(env).unwrap();
58        root.push_str(head);
59        root.push_str(&val);
60    }
61    root.push_str(copy);
62    std::fs::canonicalize(&root).unwrap_or_else(|_| panic!("{:?} is not a valid directory", root))
63}
64
65fn strip_quotes(tokens: TokenStream) -> String {
66    let mut raw = tokens.to_string();
67    if !raw.starts_with('"') || !raw.ends_with('"') || raw.len() < 3 {
68        panic!("directory must be a non-empty string");
69    }
70    raw.pop();
71    raw.remove(0);
72    raw
73}
74
75#[cfg(any(not(debug_assertions), feature = "always-embed"))]
76fn internal_dir_as_map(input: TokenStream) -> TokenStream {
77    let raw = strip_quotes(input);
78    let root = env_expand_dir(&raw);
79    let data =
80        dir_to_map(&root, &root).unwrap_or_else(|_| panic!("{:?} contains invalid entries", root));
81
82    let it = data.iter().map(|(k, v)| {
83        let u = v.iter();
84        quote! {
85            (String::from(#k), Vec::from([ #( #u ),* ]))
86        }
87    });
88
89    let output = quote! {
90        DirMap::from([ #( #it ),* ])
91    };
92
93    output.into()
94}
95
96#[cfg(all(debug_assertions, not(feature = "always-embed")))]
97fn internal_dir_as_map(input: TokenStream) -> TokenStream {
98    let raw = strip_quotes(input);
99    let root = env_expand_dir(&raw);
100    let root = root
101        .to_str()
102        .unwrap_or_else(|| panic!("{:?} cannot be converted to utf-8", root));
103
104    let output = quote! {
105        {
106            use std::{
107                collections::HashMap,
108                fs::File,
109                io::Read,
110                path::Path
111            };
112            fn file_to_bytes(path: &Path) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
113                let mut data = Vec::new();
114                let mut file = File::open(path)?;
115                file.read_to_end(&mut data)?;
116                Ok(data)
117            }
118            fn dir_to_map(root: &Path, base: &Path) -> Result<DirMap, Box<dyn std::error::Error>> {
119                let mut paths = HashMap::new();
120                for entry in std::fs::read_dir(base)? {
121                    let entry = entry?;
122                    let path = entry.path();
123                    let metadata = std::fs::metadata(&path)?;
124                    if metadata.is_file() {
125                        let data = file_to_bytes(&path)?;
126                        let rel = path.strip_prefix(root)?.to_str().unwrap();
127                        paths.insert(rel.to_string(), data);
128                    } else if metadata.is_dir() {
129                        let dirmap = dir_to_map(root, &path)?;
130                        paths.extend(dirmap);
131                    } else {
132                        panic!("{:?} is not a file or directory", entry);
133                    }
134                }
135                Ok(paths)
136            }
137            dir_to_map(Path::new(#root), Path::new(#root)).unwrap_or_else(|_| panic!("{:?} contains invalid entries", #root))
138        }
139    };
140
141    output.into()
142}
143
144/// The procedural macro magic that embeds files from a directory into the rust
145/// binary as a hashmap.
146///
147/// The input must be a string literal that represents the directory to embed.
148/// This string can contain environment variables which will be expanded at
149/// compile time.
150///
151/// The output is a `DirMap` which is an alias for `HashMap<String, Vec<u8>>`.
152/// This hashmap maps the relative path of each file to its contents as a vector
153/// of bytes.
154///
155/// By default, the files are read from the filesystem at runtime in debug mode
156/// for compilation speed. To override this behavior, enable the `always-embed`
157/// feature in `Cargo.toml`.
158///
159/// # Panics
160///
161/// This function will panic if the directory does not exist or if any file in
162/// the directory cannot be read.
163///
164/// # Examples
165///
166/// ```ignore
167/// use include_dir_as_map::{include_dir_as_map, DirMap};
168///
169/// let dirmap: DirMap = include_dir_as_map!("$CARGO_MANIFEST_DIR");
170/// let bytes = dirmap.get("Cargo.toml")?;
171/// ```
172#[proc_macro]
173pub fn include_dir_as_map(input: TokenStream) -> TokenStream {
174    internal_dir_as_map(input)
175}
176
177#[cfg(test)]
178mod tests {
179    use crate::*;
180
181    #[test]
182    fn valid_dir() {
183        let root = Path::new(".");
184        let _data = dir_to_map(root, root).unwrap();
185    }
186
187    #[test]
188    #[should_panic]
189    fn invalid_dir() {
190        let root = Path::new("invalid_path_for_testing");
191        let _data = dir_to_map(root, root).unwrap();
192    }
193
194    #[test]
195    fn valid_env_var() {
196        let path = env_expand_dir("$CARGO_MANIFEST_DIR");
197        let ok = path.ends_with("proc");
198        assert_eq!(ok, true);
199    }
200
201    #[test]
202    fn valid_env_var_2() {
203        std::env::set_var("SRC_DIR", "src");
204        let path = env_expand_dir("$CARGO_MANIFEST_DIR/$SRC_DIR");
205        let ok = path.ends_with("src");
206        assert_eq!(ok, true);
207    }
208
209    #[test]
210    #[should_panic]
211    fn invalid_env_var() {
212        let _path = env_expand_dir("$INVALID_ENV_VAR_FOR_TESTING");
213    }
214}
proc_include_dir_as_map/lib.rs

proc_include_dir_as_map/
lib.rs
