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