python_mod/
lib.rs

1use std::fs;
2use std::path::Path;
3
4use proc_macro::TokenStream;
5use proc_macro_error2::{abort, abort_call_site, proc_macro_error};
6
7use python_ast::{
8    parse_enhanced, tree::Module, CodeGen, CodeGenContext, PythonOptions, SymbolTableScopes,
9};
10
11use quote::{format_ident, quote};
12
13/// Reads the named Python module from the root of the current crate and
14/// parses it, returning the structured parse error on failure so the caller
15/// can report it against the macro invocation.
16fn load_module(mod_name: &str) -> Result<Module, python_ast::Error> {
17    let mod_name_dir = format!("src/{}/__init__.py", mod_name);
18    let mod_name_file = format!("src/{}.py", mod_name);
19
20    let (path, python_str) = if Path::new(&mod_name_dir).exists() {
21        (mod_name_dir.clone(), fs::read_to_string(&mod_name_dir))
22    } else if Path::new(&mod_name_file).exists() {
23        (mod_name_file.clone(), fs::read_to_string(&mod_name_file))
24    } else {
25        return Err(python_ast::parsing_error(
26            python_ast::SourceLocation::new(mod_name_file.clone()),
27            format!(
28                "Python module `{}` not found (looked for {} and {})",
29                mod_name, mod_name_file, mod_name_dir
30            ),
31            "create the module file, or check the module name passed to the macro",
32        ));
33    };
34
35    let python_str = python_str.map_err(|io_err| {
36        python_ast::parsing_error(
37            python_ast::SourceLocation::new(path.clone()),
38            format!("cannot read Python module file {}: {}", path, io_err),
39            "check the file's permissions and encoding (it must be valid UTF-8)",
40        )
41    })?;
42
43    parse_enhanced(&python_str, &path)
44}
45
46/// Render a structured python-ast error as a rustc diagnostic on `span_token`
47/// and abort macro expansion. The location and message are the primary line;
48/// any help text carried by the error becomes a rustc `help:` note.
49fn abort_with_python_error(span_token: &proc_macro2::TokenTree, err: &python_ast::Error) -> ! {
50    let location = err.get_field("location").unwrap_or_default().to_string();
51    let message = {
52        let m = err.get_field("message").unwrap_or_default().to_string();
53        if m.is_empty() { err.to_string() } else { m }
54    };
55    let help = err.get_field("help").unwrap_or_default().to_string();
56
57    let headline = if location.is_empty() {
58        format!("rython: {}", message)
59    } else {
60        format!("rython: {}: {}", location, message)
61    };
62
63    if help.is_empty() {
64        abort!(span_token, "{}", headline);
65    } else {
66        abort!(span_token, "{}", headline; help = "{}", help);
67    }
68}
69
70fn module_options(input: TokenStream, options: PythonOptions) -> TokenStream {
71    // We take the first token off of the stream, which is the module name, and
72    // then any other tokens will be placed at the top of the generated module.
73    let input2 = proc_macro2::TokenStream::from(input);
74    let mut tokens = input2.into_iter();
75
76    let Some(mod_name_token) = tokens.next() else {
77        abort_call_site!(
78            "rython: missing module name";
79            help = "use the macro as `python_module!(my_module)`, where `src/my_module.py` \
80                    (or `src/my_module/__init__.py`) contains the Python source"
81        );
82    };
83    let mod_name = mod_name_token.to_string();
84
85    let mut remaining_input: proc_macro2::TokenStream = tokens.collect();
86
87    // Load and parse the Python module; report failures on the module-name
88    // token so rustc points at the user's macro invocation.
89    let py_mod = match load_module(&mod_name) {
90        Ok(module) => module,
91        Err(e) => abort_with_python_error(&mod_name_token, &e),
92    };
93
94    // Collect the module's symbols before code generation (an empty symbol
95    // table changes how names resolve during codegen).
96    let symbols = py_mod.clone().find_symbols(SymbolTableScopes::new());
97
98    // Convert the parse tree to a Rust TokenStream.
99    let py_output = py_mod
100        .to_rust(CodeGenContext::Module(mod_name.clone()), options, symbols)
101        .unwrap_or_else(|e| {
102            // Recover the structured error if there is one; otherwise show
103            // the full source chain.
104            if let Some(structured) = e.downcast_ref::<python_ast::Error>() {
105                abort_with_python_error(&mod_name_token, structured)
106            } else {
107                abort!(
108                    &mod_name_token,
109                    "rython: failed to compile Python module `{}`: {}",
110                    mod_name,
111                    python_ast::format_error_chain(e.as_ref())
112                );
113            }
114        });
115
116    // Add the output to the remaining tokens, which serve as a preamble.
117    remaining_input.extend(py_output);
118
119    // Add the module declaration.
120    let new_mod_name = format_ident!("{}", mod_name);
121    let result = quote!(mod #new_mod_name {
122        #remaining_input
123    });
124
125    result.into()
126}
127
128/// Macro taking two parameters. The first is the name of the module, the second is the path to load it from.
129///
130/// ```Rust
131/// use std::module_path;
132///
133/// python_module!(py_mod);
134/// ```
135#[proc_macro]
136#[proc_macro_error]
137pub fn python_module(input: TokenStream) -> TokenStream {
138    let options = PythonOptions::default();
139    module_options(input, options)
140}
141
142#[proc_macro]
143#[proc_macro_error]
144/// Loads a python module, but does not include the stdpython libraries.
145pub fn python_module_nostd(input: TokenStream) -> TokenStream {
146    let mut options = PythonOptions::default();
147    options.with_std_python = false;
148    module_options(input, options)
149}
python_mod/lib.rs

python_mod/
lib.rs
