trillium_include_dir_impl/
lib.rs

1//! Implementation crate for the [include_dir!()] macro.
2//!
3//! [include_dir!()]: https://github.com/Michael-F-Bryan/include_dir
4
5use crate::dir::Dir;
6use proc_macro::{TokenStream, TokenTree};
7use quote::{quote, ToTokens};
8use std::{
9    borrow::Cow,
10    env,
11    path::{Path, PathBuf, StripPrefixError},
12    time::{SystemTime, UNIX_EPOCH},
13};
14
15mod dir;
16mod file;
17
18#[derive(Debug)]
19enum Error {
20    Io(std::io::Error),
21    Str(&'static str),
22    StripPrefix(StripPrefixError),
23}
24
25impl std::error::Error for Error {}
26
27impl std::fmt::Display for Error {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        match self {
30            Error::Io(i) => f.write_fmt(format_args!("{}", i)),
31            Error::Str(s) => f.write_str(s),
32            Error::StripPrefix(s) => f.write_fmt(format_args!("{}", s)),
33        }
34    }
35}
36
37impl From<StripPrefixError> for Error {
38    fn from(s: StripPrefixError) -> Self {
39        Self::StripPrefix(s)
40    }
41}
42
43impl From<std::io::Error> for Error {
44    fn from(e: std::io::Error) -> Self {
45        Self::Io(e)
46    }
47}
48
49impl From<&'static str> for Error {
50    fn from(s: &'static str) -> Self {
51        Self::Str(s)
52    }
53}
54
55#[proc_macro]
56pub fn include_dir(input: TokenStream) -> TokenStream {
57    let path = tokens_to_path(input).unwrap();
58
59    Dir::from_disk(&path, &path)
60        .expect("Couldn't load the directory")
61        .to_token_stream()
62        .into()
63}
64
65fn unwrap_string_literal(lit: &proc_macro::Literal) -> Result<String, &'static str> {
66    let mut repr = lit.to_string();
67    if !repr.starts_with('"') || !repr.ends_with('"') {
68        return Err("This macro only accepts a single, non-empty string argument");
69    }
70
71    repr.remove(0);
72    repr.pop();
73
74    Ok(repr)
75}
76
77fn match_path(tokens: Vec<TokenTree>) -> Result<String, &'static str> {
78    match tokens.as_slice() {
79        [TokenTree::Literal(lit)] => unwrap_string_literal(lit),
80
81        [TokenTree::Group(group)] => match_path(group.stream().into_iter().collect()),
82
83        _ => return Err("This macro only accepts a single, non-empty string argument".into()),
84    }
85}
86
87fn tokens_to_path(input: TokenStream) -> Result<PathBuf, Cow<'static, str>> {
88    let path = match_path(input.into_iter().collect())?;
89    let crate_root =
90        env::var("CARGO_MANIFEST_DIR").map_err(|_| "cannot find cargo manifest dir")?;
91    let path = PathBuf::from(crate_root).join(path);
92
93    if !path.exists() {
94        return Err(format!("\"{}\" doesn't exist", path.display()).into());
95    }
96
97    Ok(path
98        .canonicalize()
99        .map_err(|_| "Can't normalize the path")?)
100}
101
102#[proc_macro]
103pub fn try_include_dir(input: TokenStream) -> TokenStream {
104    let path = match tokens_to_path(input) {
105        Ok(path) => path,
106        Err(e) => return quote! { Result::<Dir, &'static str>::Err(#e) }.into(),
107    };
108
109    TokenStream::from(match load_dir(&path) {
110        Ok(dir) => quote! { Result::<Dir, &'static str>::Ok(#dir) },
111        Err(err) => quote! { Result::<Dir, &'static str>::Err(#err) },
112    })
113}
114
115fn load_dir(path: impl AsRef<Path>) -> Result<Dir, String> {
116    let path = path.as_ref();
117
118    Dir::from_disk(&path, &path)
119        .map_err(|e| format!("Couldn't load the directory: {}", e.to_string()))
120}
121
122pub(crate) fn timestamp_to_tokenstream(
123    time: std::io::Result<SystemTime>,
124) -> proc_macro2::TokenStream {
125    time.ok()
126        .and_then(|m| m.duration_since(UNIX_EPOCH).ok())
127        .map(|dur| dur.as_secs_f64())
128        .map(|secs| quote! { Some(#secs) }.to_token_stream())
129        .unwrap_or_else(|| quote! { None }.to_token_stream())
130}
trillium_include_dir_impl/lib.rs

trillium_include_dir_impl/
lib.rs
