inline_python_macros/
lib.rs

1//! Helper crate for `inline-python`.
2
3#![feature(proc_macro_span)]
4
5extern crate proc_macro;
6
7use proc_macro::{Literal, Span, TokenStream, TokenTree};
8use pyo3::{Py, Python};
9use std::{
10    collections::BTreeMap,
11    ffi::{CStr, CString},
12};
13
14mod shared;
15use shared::*;
16
17#[doc(hidden)]
18#[proc_macro]
19pub fn python(input: TokenStream) -> TokenStream {
20    python_impl(input).unwrap_or_else(|e| e)
21}
22
23#[rustfmt::skip]
24fn python_impl(input: TokenStream) -> Result<TokenStream, TokenStream> {
25    let mut variables = BTreeMap::new();
26    let python = CString::new(python_from_macro(input.clone(), Some(&mut variables))?).unwrap();
27    let filename = CString::new(Span::call_site().file()).unwrap();
28    let bytecode = compile_to_bytecode(&python, &filename, input)?;
29    Ok(TokenStream::from_iter([
30        punct(':'), punct(':'), ident("inline_python"),
31        punct(':'), punct(':'), ident("_python_block"),
32        punct('!'),
33        braces(
34            [TokenTree::Literal(bytecode)].into_iter()
35            .chain(variables.into_iter().map(|(_, value)| TokenTree::Ident(value)))
36        ),
37    ]))
38}
39
40fn compile_to_bytecode(
41    python: &CStr,
42    filename: &CStr,
43    tokens: TokenStream,
44) -> Result<Literal, TokenStream> {
45    Python::with_gil(|py| {
46        let compiled = compile_python(py, python, filename, tokens)?;
47        let bytes = unsafe {
48            let ptr =
49                pyo3::ffi::PyMarshal_WriteObjectToString(compiled.as_ptr(), pyo3::marshal::VERSION);
50            Py::from_owned_ptr(py, ptr)
51        };
52        Ok(Literal::byte_string(bytes.as_bytes(py)))
53    })
54}