Skip to content

Commit 2b433d5

Browse files
Add support for creating functions with the c-api (#7984)
1 parent 26dfeae commit 2b433d5

6 files changed

Lines changed: 379 additions & 2 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/capi/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ license.workspace = true
1212
crate-type = ["cdylib", "rlib"]
1313

1414
[dependencies]
15+
bitflags = { workspace = true }
1516
num-complex = { workspace = true }
1617
rustpython-vm = { workspace = true, features = ["threading", "compiler"] }
1718
rustpython-stdlib = {workspace = true, features = ["threading"] }

‎crates/capi/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ pub mod floatobject;
1818
pub mod import;
1919
pub mod listobject;
2020
pub mod longobject;
21+
pub mod methodobject;
2122
pub mod object;
2223
pub mod pycapsule;
2324
pub mod pyerrors;

‎crates/capi/src/methodobject.rs‎

Lines changed: 369 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,369 @@
1+
use crate::PyObject;
2+
use crate::object::PyTypeObject;
3+
use crate::object::define_py_check;
4+
use crate::pystate::with_vm;
5+
use core::ffi::{CStr, c_char, c_int};
6+
use core::ptr::NonNull;
7+
use rustpython_vm::function::{FuncArgs, HeapMethodDef, PosArgs, PyMethodFlags};
8+
use rustpython_vm::{AsObject, PyObjectRef, PyRef, PyResult, VirtualMachine};
9+
10+
define_py_check!(fn PyCFunction_Check, types.builtin_function_or_method_type);
11+
define_py_check!(exact fn PyCFunction_CheckExact, types.builtin_function_or_method_type);
12+
13+
#[repr(C)]
14+
pub struct PyMethodDef {
15+
pub ml_name: *const c_char,
16+
pub ml_meth: PyMethodPointer,
17+
pub ml_flags: c_int,
18+
pub ml_doc: *const c_char,
19+
}
20+
21+
#[repr(C)]
22+
#[derive(Copy, Clone)]
23+
#[allow(non_snake_case)]
24+
pub union PyMethodPointer {
25+
pub PyCFunction: unsafe extern "C" fn(slf: *mut PyObject, args: *mut PyObject) -> *mut PyObject,
26+
pub PyCFunctionWithKeywords: unsafe extern "C" fn(
27+
slf: *mut PyObject,
28+
args: *mut PyObject,
29+
kwargs: *mut PyObject,
30+
) -> *mut PyObject,
31+
pub PyCFunctionFast: unsafe extern "C" fn(
32+
slf: *mut PyObject,
33+
args: *const *mut PyObject,
34+
nargs: isize,
35+
) -> *mut PyObject,
36+
pub PyCFunctionFastWithKeywords: unsafe extern "C" fn(
37+
slf: *mut PyObject,
38+
args: *const *mut PyObject,
39+
nargs: isize,
40+
kwnames: *mut PyObject,
41+
) -> *mut PyObject,
42+
}
43+
44+
pub(crate) fn build_method_def(
45+
vm: &VirtualMachine,
46+
ml: &PyMethodDef,
47+
has_self: bool,
48+
) -> PyResult<PyRef<HeapMethodDef>> {
49+
let name = unsafe { CStr::from_ptr(ml.ml_name) }
50+
.to_str()
51+
.map_err(|_| vm.new_system_error("Method name was not valid UTF-8"))?;
52+
53+
let doc = NonNull::new(ml.ml_doc.cast_mut())
54+
.map(|doc| {
55+
unsafe { CStr::from_ptr(doc.as_ptr()) }
56+
.to_str()
57+
.map_err(|_| vm.new_system_error("Method doc was not valid UTF-8"))
58+
})
59+
.transpose()?;
60+
61+
let flags = PyMethodFlags::from_bits(ml.ml_flags as u32)
62+
.ok_or_else(|| vm.new_system_error("PyMethodDef contains unknown flags"))?;
63+
64+
let method = ml.ml_meth;
65+
66+
if flags.contains(PyMethodFlags::METHOD) {
67+
return Err(vm.new_system_error("METH_METHOD is not supported on abi3"));
68+
}
69+
70+
let call_flags = flags
71+
& (PyMethodFlags::VARARGS
72+
| PyMethodFlags::KEYWORDS
73+
| PyMethodFlags::NOARGS
74+
| PyMethodFlags::O
75+
| PyMethodFlags::FASTCALL);
76+
77+
bitflags::bitflags_match!(call_flags, {
78+
PyMethodFlags::NOARGS => {
79+
if has_self {
80+
let callable = move |zelf: PyObjectRef, vm: &VirtualMachine| unsafe {
81+
let f = method.PyCFunction;
82+
let ret_ptr = f(zelf.as_raw().cast_mut(), core::ptr::null_mut());
83+
ret_ptr_to_pyresult(vm, ret_ptr)
84+
};
85+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
86+
} else {
87+
let callable = move |vm: &VirtualMachine| unsafe {
88+
let f = method.PyCFunction;
89+
let ret_ptr = f(core::ptr::null_mut(), core::ptr::null_mut());
90+
ret_ptr_to_pyresult(vm, ret_ptr)
91+
};
92+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
93+
}
94+
},
95+
PyMethodFlags::VARARGS => {
96+
let callable = move |args: PosArgs, vm: &VirtualMachine| unsafe {
97+
call_function(vm, method, flags, Some(args))
98+
};
99+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
100+
},
101+
PyMethodFlags::VARARGS | PyMethodFlags::KEYWORDS => {
102+
let callable = move | args: FuncArgs, vm: &VirtualMachine| unsafe {
103+
call_function_with_keywords(vm, method, flags, args)
104+
};
105+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
106+
},
107+
PyMethodFlags::FASTCALL | PyMethodFlags::KEYWORDS => {
108+
let callable = move |args: FuncArgs, vm: &VirtualMachine| unsafe {
109+
call_fast_function_with_keywords(vm, method, flags, args)
110+
};
111+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
112+
},
113+
PyMethodFlags::FASTCALL => {
114+
let callable = move |args: PosArgs, vm: &VirtualMachine| unsafe {
115+
call_fast_function(vm, method, flags, args)
116+
};
117+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
118+
},
119+
PyMethodFlags::O => {
120+
let f = unsafe { method.PyCFunction };
121+
if has_self {
122+
let callable = move |zelf: PyObjectRef, arg: PyObjectRef, vm: &VirtualMachine| -> PyResult {
123+
let ret_ptr = unsafe { f(zelf.as_raw().cast_mut(), arg.as_raw().cast_mut()) };
124+
ret_ptr_to_pyresult(vm, ret_ptr)
125+
};
126+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
127+
} else {
128+
let callable = move |arg: PyObjectRef, vm: &VirtualMachine| -> PyResult {
129+
let ret_ptr = unsafe { f(core::ptr::null_mut(), arg.as_raw().cast_mut()) };
130+
ret_ptr_to_pyresult(vm, ret_ptr)
131+
};
132+
Ok(vm.ctx.new_method_def(name, callable, flags, doc))
133+
}
134+
},
135+
_ => {
136+
Err(vm.new_system_error(format!(
137+
"function {name} has unsupported or invalid calling-convention flags: {flags:?}"
138+
)))
139+
},
140+
})
141+
}
142+
143+
unsafe fn call_function<A: Into<FuncArgs>>(
144+
vm: &VirtualMachine,
145+
method: PyMethodPointer,
146+
flags: PyMethodFlags,
147+
args: Option<A>,
148+
) -> PyResult {
149+
let f = unsafe { method.PyCFunction };
150+
let (slf, arg_tuple) = if let Some(mut args) = args.map(Into::into) {
151+
let slf = take_self_arg(&mut args, flags);
152+
let arg_tuple = vm.ctx.new_tuple(args.args);
153+
(slf, Some(arg_tuple))
154+
} else {
155+
(None, None)
156+
};
157+
158+
let slf_ptr = slf
159+
.as_ref()
160+
.map(|obj| obj.as_object().as_raw().cast_mut())
161+
.unwrap_or_default();
162+
163+
let arg_ptr = arg_tuple
164+
.as_ref()
165+
.map(|tuple| tuple.as_object().as_raw().cast_mut())
166+
.unwrap_or_default();
167+
168+
let ret_ptr = unsafe { f(slf_ptr, arg_ptr) };
169+
ret_ptr_to_pyresult(vm, ret_ptr)
170+
}
171+
172+
unsafe fn call_function_with_keywords(
173+
vm: &VirtualMachine,
174+
method: PyMethodPointer,
175+
flags: PyMethodFlags,
176+
mut args: FuncArgs,
177+
) -> PyResult {
178+
let f = unsafe { method.PyCFunctionWithKeywords };
179+
let slf = take_self_arg(&mut args, flags);
180+
let slf_ptr = slf
181+
.as_ref()
182+
.map(|obj| obj.as_object().as_raw().cast_mut())
183+
.unwrap_or_default();
184+
let arg_tuple = vm.ctx.new_tuple(args.args);
185+
let kwargs = vm.ctx.new_dict();
186+
for (k, v) in args.kwargs {
187+
kwargs.set_item(&*k, v, vm)?;
188+
}
189+
let ret_ptr = unsafe {
190+
f(
191+
slf_ptr,
192+
arg_tuple.as_object().as_raw().cast_mut(),
193+
kwargs.as_object().as_raw().cast_mut(),
194+
)
195+
};
196+
ret_ptr_to_pyresult(vm, ret_ptr)
197+
}
198+
199+
unsafe fn call_fast_function_with_keywords(
200+
vm: &VirtualMachine,
201+
method: PyMethodPointer,
202+
flags: PyMethodFlags,
203+
mut args: FuncArgs,
204+
) -> PyResult {
205+
let f = unsafe { method.PyCFunctionFastWithKeywords };
206+
let slf = take_self_arg(&mut args, flags);
207+
let slf_ptr = slf
208+
.as_ref()
209+
.map(|obj| obj.as_object().as_raw().cast_mut())
210+
.unwrap_or_default();
211+
let nargs = args.args.len();
212+
let mut fastcall_args = args.args;
213+
let kwnames_tuple = if !args.kwargs.is_empty() {
214+
let mut kwnames = Vec::with_capacity(args.kwargs.len());
215+
for (k, v) in args.kwargs {
216+
kwnames.push(vm.ctx.new_str(k).into());
217+
fastcall_args.push(v);
218+
}
219+
Some(vm.ctx.new_tuple(kwnames))
220+
} else {
221+
None
222+
};
223+
let kwnames_ptr = kwnames_tuple
224+
.as_ref()
225+
.map(|tuple| tuple.as_object().as_raw().cast_mut())
226+
.unwrap_or_default();
227+
// SAFETY: PyObjectRef is repr(transparent) over a pointer to PyObject, so a
228+
// Vec<PyObjectRef> has a layout-compatible contiguous backing buffer. The
229+
// vector is kept alive for the duration of the call.
230+
let fastcall_arg_ptrs = fastcall_args.as_ptr().cast::<*mut PyObject>();
231+
let ret_ptr = unsafe { f(slf_ptr, fastcall_arg_ptrs, nargs as isize, kwnames_ptr) };
232+
ret_ptr_to_pyresult(vm, ret_ptr)
233+
}
234+
235+
unsafe fn call_fast_function(
236+
vm: &VirtualMachine,
237+
method: PyMethodPointer,
238+
flags: PyMethodFlags,
239+
args: PosArgs,
240+
) -> PyResult {
241+
let f = unsafe { method.PyCFunctionFast };
242+
let mut args: FuncArgs = args.into();
243+
let slf = take_self_arg(&mut args, flags);
244+
let slf_ptr = slf
245+
.as_ref()
246+
.map(|obj| obj.as_object().as_raw().cast_mut())
247+
.unwrap_or_default();
248+
// SAFETY: PyObjectRef is repr(transparent) over a pointer to PyObject, so a
249+
// Vec<PyObjectRef> has a layout-compatible contiguous backing buffer. The
250+
// vector is kept alive for the duration of the call.
251+
let fastcall_arg_ptrs = args.args.as_mut_ptr().cast::<*mut PyObject>();
252+
let ret_ptr = unsafe { f(slf_ptr, fastcall_arg_ptrs, args.args.len() as isize) };
253+
ret_ptr_to_pyresult(vm, ret_ptr)
254+
}
255+
256+
fn ret_ptr_to_pyresult(vm: &VirtualMachine, ret_ptr: *mut PyObject) -> PyResult {
257+
let ret_ptr = NonNull::new(ret_ptr).ok_or_else(|| {
258+
vm.take_raised_exception()
259+
.expect("Native function returned NULL, but there was no exception set")
260+
})?;
261+
Ok(unsafe { PyObjectRef::from_raw(ret_ptr) })
262+
}
263+
264+
fn take_self_arg(args: &mut FuncArgs, flags: PyMethodFlags) -> Option<PyObjectRef> {
265+
if flags.contains(PyMethodFlags::STATIC) {
266+
None
267+
} else {
268+
args.take_positional()
269+
}
270+
}
271+
272+
#[unsafe(no_mangle)]
273+
pub unsafe extern "C" fn PyCMethod_New(
274+
ml: *mut PyMethodDef,
275+
slf: *mut PyObject,
276+
_module: *mut PyObject,
277+
_cls: *mut PyTypeObject,
278+
) -> *mut PyObject {
279+
with_vm(|vm| -> PyResult {
280+
assert!(
281+
_cls.is_null(),
282+
"PyCMethod_New does not support METH_METHOD on abi3"
283+
);
284+
let ml = unsafe { &*ml };
285+
let zelf = unsafe { slf.as_ref().map(|obj| obj.to_owned()) };
286+
Ok(build_method_def(vm, ml, zelf.is_some())?
287+
.build_function(vm, zelf)
288+
.into())
289+
})
290+
}
291+
292+
#[unsafe(no_mangle)]
293+
pub unsafe extern "C" fn PyCFunction_New(
294+
ml: *mut PyMethodDef,
295+
slf: *mut PyObject,
296+
) -> *mut PyObject {
297+
unsafe { PyCMethod_New(ml, slf, core::ptr::null_mut(), core::ptr::null_mut()) }
298+
}
299+
300+
#[unsafe(no_mangle)]
301+
pub unsafe extern "C" fn PyCFunction_NewEx(
302+
ml: *mut PyMethodDef,
303+
slf: *mut PyObject,
304+
module: *mut PyObject,
305+
) -> *mut PyObject {
306+
unsafe { PyCMethod_New(ml, slf, module, core::ptr::null_mut()) }
307+
}
308+
309+
#[cfg(false)]
310+
mod tests {
311+
use pyo3::exceptions::PyException;
312+
use pyo3::ffi::{PyLong_FromLong, PyObject};
313+
use pyo3::prelude::*;
314+
use pyo3::types::{PyCFunction, PyInt, PyString};
315+
316+
#[test]
317+
fn test_closure_function() {
318+
Python::attach(|py| {
319+
let f = PyCFunction::new_closure(py, None, None, |_args, _kwargs| "Hello from Rust!")
320+
.unwrap();
321+
322+
assert_eq!(
323+
f.call0().unwrap().cast::<PyString>().unwrap(),
324+
"Hello from Rust!"
325+
);
326+
})
327+
}
328+
329+
#[test]
330+
fn test_function_no_args() {
331+
Python::attach(|py| {
332+
unsafe extern "C" fn c_fn(_self: *mut PyObject, _args: *mut PyObject) -> *mut PyObject {
333+
assert!(_self.is_null());
334+
assert!(_args.is_null());
335+
unsafe { PyLong_FromLong(4200) }
336+
}
337+
338+
let py_fn = PyCFunction::new(py, c_fn, c"py_fn", c"", None).unwrap();
339+
340+
let result = py_fn
341+
.call0()
342+
.unwrap()
343+
.cast::<PyInt>()
344+
.unwrap()
345+
.extract::<u32>()
346+
.unwrap();
347+
assert_eq!(result, 4200);
348+
349+
assert!(py_fn.call((1,), None).is_err());
350+
assert!(py_fn.call((1, 2), None).is_err());
351+
})
352+
}
353+
354+
#[test]
355+
fn test_closure_function_error() {
356+
Python::attach(|py| {
357+
let f = PyCFunction::new_closure(py, None, None, |_args, _kwargs| {
358+
Err::<(), _>(PyException::new_err("Something went wrong"))
359+
})
360+
.unwrap();
361+
362+
let err = f.call0().unwrap_err();
363+
assert_eq!(
364+
err.value(py).repr().unwrap(),
365+
"Exception('Something went wrong')"
366+
);
367+
})
368+
}
369+
}

0 commit comments

Comments
 (0)