python_mod/
lib.rs

1#![feature(proc_macro_span)]
2use std::path::Path;
3use std::fs;
4
5use proc_macro::{ Span, TokenStream  };
6use proc_macro_error::{
7    proc_macro_error,
8    abort,
9};
10
11use python_ast::{
12    CodeGen, CodeGenContext,
13    PythonOptions, PyResult,
14    parse,
15    tree::Module,
16    SymbolTableScopes
17};
18
19use quote::{quote, format_ident};
20
21/// Reads the named Python module from the root of the current crate.
22fn load_module(mod_name: String) -> PyResult<Module> {
23    let mod_path = Span::mixed_site().source().file();
24
25    let mod_name_dir = format!("src/{}/__init__.py", &mod_name);
26    let mod_name_file = format!("src/{}.py", &mod_name);
27
28    let python_str = if Path::new(&mod_name_dir).exists() {
29        fs::read_to_string(mod_name_dir)
30    } else if Path::new(&mod_name_file).exists() {
31        fs::read_to_string(mod_name_file)
32    } else {
33        abort!(mod_name, "Module not found {} in {:?}", mod_name, mod_path)
34    }.expect("Reading python module");
35
36    parse(&python_str, &mod_name)
37}
38
39fn module_options(input: TokenStream, options: PythonOptions) -> TokenStream {
40    // We take the first token off of the stream, which is the module name, and then any other tokens
41    // will be placed at the top of the generated module.
42    let symbols = SymbolTableScopes::new();
43
44    let mut new_input = input.into_iter();
45    let mod_name = new_input.next()
46        .expect("missing module name").to_string();
47
48    let remaining_input = proc_macro::TokenStream::from_iter(new_input);
49    let mut remaining_input2 = proc_macro2::TokenStream::from(remaining_input);
50
51    //let mut output = quote!(#remaining_input);
52
53    // Loads a Python module and stores it in a String
54    let py_mod = load_module(mod_name.clone())
55        .expect("loading python module ");
56
57    // Convert the parse tree to Rust a TokenStream.
58    let py_output = py_mod.to_rust(CodeGenContext::Module(mod_name.clone()), options, symbols)
59        .expect("converting Python to Rust");
60
61    // Add the output to the remaining tokens, which serve as a preamble.
62    proc_macro2::TokenStream::extend::<proc_macro2::TokenStream>(&mut remaining_input2, py_output.into());
63
64
65    // Add the module declaration.
66    let new_mod_name = format_ident!("{}", mod_name);
67    let result = quote!(mod #new_mod_name {
68        #remaining_input2
69    });
70
71    result.into()
72}
73
74/// Macro taking two parameters. The first is the name of the module, the second is the path to load it from.
75///
76/// ```Rust
77/// use std::module_path;
78///
79/// python_module!(py_mod);
80/// ```
81#[proc_macro]
82#[proc_macro_error]
83pub fn python_module(input: TokenStream) -> TokenStream {
84    let options = PythonOptions::default();
85    module_options(input, options)
86
87}
88
89#[proc_macro]
90#[proc_macro_error]
91/// Loads a python module, but does not include the stdpython libraries.
92pub fn python_module_nostd(input: TokenStream) -> TokenStream {
93    let mut options = PythonOptions::default();
94    options.with_std_python = false;
95    module_options(input, options)
96}