1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
//! Inline Python code directly in your Rust code.
//!
//! # Example
//!
//! ```
//! #![feature(proc_macro_hygiene)]
//! use inline_python::python;
//!
//! fn main() {
//!     let who = "world";
//!     let n = 5;
//!     python! {
//!         for i in range('n):
//!             print(i, "Hello", 'who)
//!         print("Goodbye")
//!     }
//! }
//! ```
//!
//! # How to use
//!
//! Use the `python!{..}` macro to write Python code directly in your Rust code.
//! You'll need to add `#![feature(proc_macro_hygiene)]`, and use a nightly
//! version of the compiler that supports this feature.
//!
//! ## Using Rust variables
//!
//! To reference Rust variables, use `'var`, as shown in the example above.
//! `var` needs to implement [`pyo3::ToPyObject`].
//!
//! ## Re-using a Python context
//! It is possible to create a [`Context`] object ahead of time and use it for running the Python code.
//! The context can be re-used for multiple invocations to share global variables across macro calls.
//!
//! ```
//! # #![feature(proc_macro_hygiene)]
//! # use inline_python::python;
//! let c = inline_python::Context::new();
//! python! {
//!   #![context = &c]
//!   foo = 5
//! }
//! python! {
//!   #![context = &c]
//!   assert foo == 5
//! }
//! ```
//!
//! ## Getting information back
//!
//! A [`Context`] object could also be used to pass information back to Rust,
//! as you can retrieve the global Python variables from the context through
//! [`Context::get_global`].
//!
//! ## Syntax issues
//!
//! Since the Rust tokenizer will tokenize the Python code, some valid Python
//! code is rejected. The two main things to remember are:
//!
//! - Use double quoted strings (`""`) instead of single quoted strings (`''`).
//!
//!   (Single quoted strings only work if they contain a single character, since
//!   in Rust, `'a'` is a character literal.)
//!
//! - Use `//`-comments instead of `#`-comments.
//!
//!   (If you use `#` comments, the Rust tokenizer will try to tokenize your
//!   comment, and complain if your comment doesn't tokenize properly.)
//!
//! Other minor things that don't work are:
//!
//! - Certain escape codes in string literals.
//!   (Specifically: `\a`, `\b`, `\f`, `\v`, `\N{..}`, `\123` (octal escape
//!   codes), `\u`, and `\U`.)
//!
//!   These, however, are accepted just fine: `\\`, `\n`, `\t`, `\r`, `\xAB`
//!   (hex escape codes), and `\0`
//!
//! - Raw string literals with escaped double quotes. (E.g. `r"...\"..."`.)
//!
//! - Triple-quoted byte- and raw-strings with content that would not be valid
//!   as a regular string. And the same for raw-byte and raw-format strings.
//!   (E.g. `b"""\xFF"""`, `r"""\z"""`, `fr"\z"`, `br"\xFF"`.)
//!
//! - The `//` and `//=` operators are unusable, as they start a comment.
//!
//!   Workaround: you can write `##` instead, which is automatically converted
//!   to `//`.
//!
//! Everything else should work fine.

use std::os::raw::c_char;

pub use inline_python_macros::python;
pub use pyo3;

use pyo3::{
	ffi,
	types::{PyAny, PyDict},
	AsPyPointer, FromPyObject, IntoPy, PyErr, PyObject, PyResult, Python,
};

#[doc(hidden)]
pub use std::ffi::CStr;

/// An execution context for Python code.
///
/// If you pass a manually created context to the `python!{}` macro, you can share it across invocations.
/// This will keep all global variables and imports intact between macro invocations.
///
/// ```
/// # #![feature(proc_macro_hygiene)]
/// # use inline_python::python;
/// let c = inline_python::Context::new();
/// python! {
///   #![context = &c]
///   foo = 5
/// }
/// python! {
///   #![context = &c]
///   assert foo == 5
/// }
/// ```
///
/// You may also use it to inspect global variables after the execution of the Python code.
/// Note that you need to acquire the GIL in order to access those globals:
///
/// ```
/// # #![feature(proc_macro_hygiene)]
/// use inline_python::python;
/// let context = inline_python::Context::new();
/// python! {
///   #![context = &context]
///   foo = 5
/// }
///
/// let foo: Option<i32> = context.get_global("foo").unwrap();
/// assert_eq!(foo, Some(5));
/// ```
pub struct Context {
	globals: PyObject,
}

impl Context {
	/// Create a new context for running python code.
	///
	/// This function temporarily acquires the GIL.
	/// If you already have the GIL, use [`Context::new_with_gil`] instead.
	///
	/// This function panics if it fails to create the context.
	/// See [`Context::new_checked`] for a version that returns a result.
	pub fn new() -> Self {
		let gil = Python::acquire_gil();
		let py = gil.python();
		match Self::new_with_gil(py) {
			Ok(x) => x,
			Err(error) => {
				error.print(py);
				panic!("failed to create python context");
			}
		}
	}

	/// Create a new context for running python code.
	///
	/// This function temporarily acquires the GIL.
	/// If you already have the GIL, use [`Context::new_with_gil`] instead.
	pub fn new_checked() -> PyResult<Self> {
		let gil = Python::acquire_gil();
		let py = gil.python();
		Self::new_with_gil(py)
	}

	/// Create a new context for running Python code.
	///
	/// You must acquire the GIL to call this function.
	pub fn new_with_gil(py: Python) -> PyResult<Self> {
		let main_mod = unsafe { ffi::PyImport_AddModule("__main__\0".as_ptr() as *const _) };
		if main_mod.is_null() {
			return Err(PyErr::fetch(py));
		};

		let globals = PyDict::new(py);
		if unsafe { ffi::PyDict_Merge(globals.as_ptr(), ffi::PyModule_GetDict(main_mod), 0) != 0 } {
			return Err(PyErr::fetch(py));
		}

		Ok(Self {
			globals: globals.into_py(py),
		})
	}

	/// Get the globals as dictionary.
	pub fn globals<'p>(&self, py: Python<'p>) -> &'p PyDict {
		unsafe { py.from_borrowed_ptr(self.globals.as_ptr()) }
	}

	/// Retrieve a global variable from the context.
	///
	/// This function temporarily acquires the GIL.
	/// If you already have the GIL, use [`Context::get_global_with_gil`] instead.
	pub fn get_global<T: for<'p> FromPyObject<'p>>(&self, name: &str) -> PyResult<Option<T>> {
		self.get_global_with_gil(Python::acquire_gil().python(), name)
	}

	/// Retrieve a global variable from the context.
	pub fn get_global_with_gil<'p, T: FromPyObject<'p>>(&self, py: Python<'p>, name: &str) -> PyResult<Option<T>> {
		match self.globals(py).get_item(name) {
			None => Ok(None),
			Some(value) => FromPyObject::extract(value).map(Some),
		}
	}
}

#[doc(hidden)]
pub fn run_python_code<'p>(py: Python<'p>, context: &Context, compiled_code: &[u8], rust_vars: Option<&PyDict>) -> PyResult<&'p PyAny> {
	unsafe {
		// Add the rust variable in a global dictionary named RUST.
		// If no rust vars are given, make the RUST global an empty dictionary.
		let rust_vars = rust_vars.unwrap_or_else(|| PyDict::new(py)).as_ptr();
		if ffi::PyDict_SetItemString(context.globals.as_ptr(), "RUST\0".as_ptr() as *const _, rust_vars) != 0 {
			return Err(PyErr::fetch(py));
		}

		let compiled_code = python_unmarshal_object_from_bytes(py, compiled_code)?;
		let result = ffi::PyEval_EvalCode(compiled_code.as_ptr(), context.globals.as_ptr(), std::ptr::null_mut());

		py.from_owned_ptr_or_err(result)
	}
}

extern "C" {
	fn PyMarshal_ReadObjectFromString(data: *const c_char, len: isize) -> *mut ffi::PyObject;
}

/// Use built-in python marshal support to read an object from bytes.
fn python_unmarshal_object_from_bytes(py: Python, data: &[u8]) -> pyo3::PyResult<PyObject> {
	unsafe {
		let object = PyMarshal_ReadObjectFromString(data.as_ptr() as *const c_char, data.len() as isize);
		if object.is_null() {
			return Err(PyErr::fetch(py));
		}

		Ok(PyObject::from_owned_ptr(py, object))
	}
}