script_macro/
lib.rs

1#![doc = include_str!("../README.md")]
2
3extern crate proc_macro;
4
5use std::path::PathBuf;
6
7use proc_macro::TokenStream;
8use rhai::{Engine, EvalAltResult, ImmutableString, Position, Scope};
9use syn::{
10    parse::{Parse, ParseStream},
11    parse_macro_input, LitStr,
12};
13
14struct RunScriptInput {
15    script_source: LitStr,
16}
17
18impl Parse for RunScriptInput {
19    fn parse(input: ParseStream) -> syn::Result<Self> {
20        Ok(Self {
21            script_source: input.parse()?,
22        })
23    }
24}
25
26fn get_source_context(source_code: &str, padding: usize, pos: Position) -> String {
27    let mut source_snippet = String::new();
28
29    if let Some(lineno) = pos.line() {
30        let lines: Vec<_> = source_code.split('\n').collect();
31        for (i, line) in lines
32            [(lineno - padding).clamp(0, lines.len())..(lineno + padding).clamp(0, lines.len())]
33            .iter()
34            .enumerate()
35        {
36            if i == padding - 1 {
37                source_snippet.push_str("--> ");
38            } else {
39                source_snippet.push_str("    ");
40            }
41
42            source_snippet.push_str(line);
43            source_snippet.push('\n');
44        }
45    }
46
47    source_snippet
48}
49
50fn handle_runtime_error(source_code: &str, e: Box<EvalAltResult>) {
51    let pos = {
52        let mut inner_error = &e;
53
54        while let EvalAltResult::ErrorInModule(_, err, ..)
55        | EvalAltResult::ErrorInFunctionCall(_, _, err, ..) = &**inner_error
56        {
57            inner_error = err;
58        }
59
60        inner_error.position()
61    };
62
63    panic!("{}\n\n{}", e, get_source_context(source_code, 3, pos));
64}
65
66#[proc_macro]
67pub fn run_script(params: TokenStream) -> TokenStream {
68    let args = parse_macro_input!(params as RunScriptInput);
69
70    let engine = get_default_engine();
71    let output: String = engine
72        .eval(&args.script_source.value())
73        .map_err(|e| handle_runtime_error(&args.script_source.value(), e))
74        .unwrap();
75
76    output.parse().expect("invalid token stream")
77}
78
79#[proc_macro_attribute]
80pub fn run_script_on(params: TokenStream, item: TokenStream) -> TokenStream {
81    let args = parse_macro_input!(params as RunScriptInput);
82    let engine = get_default_engine();
83
84    let mut scope = Scope::new();
85    scope.push("item", item.to_string());
86    let output: String = engine
87        .eval_with_scope(&mut scope, &args.script_source.value())
88        .map_err(|e| handle_runtime_error(&args.script_source.value(), e))
89        .unwrap();
90
91    output.parse().expect("invalid token stream")
92}
93
94fn get_default_engine() -> Engine {
95    let mut engine = Engine::new();
96
97    engine.set_max_expr_depths(100, 100);
98
99    #[cfg(feature = "parse-yaml")]
100    engine.register_fn("parse_yaml", helper_parse_yaml);
101    #[cfg(feature = "parse-json")]
102    engine.register_fn("parse_json", helper_parse_json);
103    #[cfg(feature = "parse-yaml")]
104    engine.register_fn("stringify_yaml", helper_stringify_yaml);
105    #[cfg(feature = "parse-json")]
106    engine.register_fn("stringify_json", helper_stringify_json);
107    engine.register_fn("slugify_ident", helper_slugify_ident);
108    #[cfg(feature = "glob")]
109    engine.register_fn("glob", helper_glob);
110    engine.register_fn("basename", helper_basename);
111
112    #[cfg(feature = "filesystem")]
113    {
114        use rhai::packages::Package;
115        use rhai_fs::FilesystemPackage;
116        let package = FilesystemPackage::new();
117        package.register_into_engine(&mut engine);
118    }
119
120    engine
121}
122
123#[cfg(any(
124    feature = "parse-yaml",
125    feature = "parse-json",
126    feature = "filesystem",
127    feature = "glob",
128))]
129fn coerce_err(x: impl std::fmt::Debug) -> Box<EvalAltResult> {
130    format!("{x:?}").into()
131}
132
133#[cfg(feature = "parse-yaml")]
134fn helper_parse_yaml(input: ImmutableString) -> Result<rhai::Dynamic, Box<EvalAltResult>> {
135    serde_yaml::from_str(input.as_str()).map_err(coerce_err)
136}
137
138#[cfg(feature = "parse-yaml")]
139fn helper_stringify_yaml(input: rhai::Dynamic) -> Result<ImmutableString, Box<EvalAltResult>> {
140    serde_yaml::to_string(&input)
141        .map(From::from)
142        .map_err(coerce_err)
143}
144
145#[cfg(feature = "parse-json")]
146fn helper_parse_json(input: ImmutableString) -> Result<rhai::Dynamic, Box<EvalAltResult>> {
147    serde_json::from_str(input.as_str()).map_err(coerce_err)
148}
149
150#[cfg(feature = "parse-json")]
151fn helper_stringify_json(input: rhai::Dynamic) -> Result<ImmutableString, Box<EvalAltResult>> {
152    serde_json::to_string(&input)
153        .map(From::from)
154        .map_err(coerce_err)
155}
156
157fn helper_slugify_ident(input: ImmutableString) -> ImmutableString {
158    let mut is_first_char = true;
159    input
160        .as_str()
161        .replace(
162            |x: char| {
163                if is_first_char && x.is_ascii_digit() {
164                    return true;
165                }
166                is_first_char = false;
167
168                !matches!(x, 'a'..='z' | 'A'..='Z' | '0'..='9' | '_')
169            },
170            "_",
171        )
172        .into()
173}
174
175#[cfg(feature = "glob")]
176fn helper_glob(pattern: ImmutableString) -> Result<rhai::Dynamic, Box<EvalAltResult>> {
177    let mut result = Vec::new();
178
179    for entry in glob::glob(pattern.as_str()).map_err(coerce_err)? {
180        let entry = entry.map_err(coerce_err)?;
181
182        result.push(entry);
183    }
184
185    Ok(result.into())
186}
187
188fn helper_basename(input: PathBuf) -> Result<ImmutableString, Box<EvalAltResult>> {
189    Ok(input
190        .file_name()
191        .unwrap_or(input.as_os_str())
192        .to_str()
193        .ok_or("basename is not valid unicode")?
194        .to_owned()
195        .into())
196}
script_macro/lib.rs

script_macro/
lib.rs
