inline_python_macros/
lib.rs

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