diff --git a/src/rust/engine/src/externs/interface.rs b/src/rust/engine/src/externs/interface.rs index a75bc027a952..3bb073b2831f 100644 --- a/src/rust/engine/src/externs/interface.rs +++ b/src/rust/engine/src/externs/interface.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use std::time::Duration; use async_latch::AsyncLatch; -use cpython::{PyClone, PyResult as CPythonPyResult, ToPyObject}; +use cpython::{PyResult as CPythonPyResult, ToPyObject}; use futures::future::FutureExt; use futures::future::{self, TryFutureExt}; use futures::Future; @@ -33,7 +33,8 @@ use petgraph::graph::{DiGraph, Graph}; use process_execution::RemoteCacheWarningsBehavior; use pyo3::exceptions::{PyException, PyValueError}; use pyo3::prelude::{ - pyclass, pymethods, pymodule, Py, PyModule, PyObject, PyResult as PyO3Result, Python, + pyclass, pyfunction, pymethods, pymodule, wrap_pyfunction, Py, PyModule, PyObject, + PyResult as PyO3Result, Python, }; use pyo3::types::{PyList, PyString, PyTuple, PyType}; use regex::Regex; @@ -72,6 +73,30 @@ fn native_engine_new(py: Python, m: &PyModule) -> PyO3Result<()> { m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(stdio_initialize, m)?)?; + m.add_function(wrap_pyfunction!(stdio_thread_console_set, m)?)?; + m.add_function(wrap_pyfunction!(stdio_thread_console_color_mode_set, m)?)?; + m.add_function(wrap_pyfunction!(stdio_thread_console_clear, m)?)?; + m.add_function(wrap_pyfunction!(stdio_thread_get_destination, m)?)?; + m.add_function(wrap_pyfunction!(stdio_thread_set_destination, m)?)?; + + m.add_function(wrap_pyfunction!(flush_log, m)?)?; + m.add_function(wrap_pyfunction!(write_log, m)?)?; + m.add_function(wrap_pyfunction!(set_per_run_log_path, m)?)?; + m.add_function(wrap_pyfunction!(teardown_dynamic_ui, m)?)?; + m.add_function(wrap_pyfunction!(maybe_set_panic_handler, m)?)?; + + m.add_function(wrap_pyfunction!(task_side_effected, m)?)?; + + m.add_function(wrap_pyfunction!(tasks_task_begin, m)?)?; + m.add_function(wrap_pyfunction!(tasks_task_end, m)?)?; + m.add_function(wrap_pyfunction!(tasks_add_get, m)?)?; + m.add_function(wrap_pyfunction!(tasks_add_union, m)?)?; + m.add_function(wrap_pyfunction!(tasks_add_select, m)?)?; + m.add_function(wrap_pyfunction!(tasks_add_query, m)?)?; + + m.add_function(wrap_pyfunction!(strongly_connected_components, m)?)?; + Ok(()) } @@ -81,80 +106,6 @@ cpython::py_module_initializer!(native_engine, |py, m| { m.add(py, "PollTimeout", py.get_type::()) .unwrap(); - m.add( - py, - "stdio_initialize", - cpython::py_fn!( - py, - stdio_initialize( - a: u64, - b: bool, - d: bool, - e: cpython::PyDict, - f: Vec, - g: Vec, - h: String - ) - ), - )?; - m.add( - py, - "stdio_thread_console_set", - cpython::py_fn!( - py, - stdio_thread_console_set(stdin_fileno: i32, stdout_fileno: i32, stderr_fileno: i32) - ), - )?; - m.add( - py, - "stdio_thread_console_color_mode_set", - cpython::py_fn!(py, stdio_thread_console_color_mode_set(use_color: bool)), - )?; - m.add( - py, - "stdio_thread_console_clear", - cpython::py_fn!(py, stdio_thread_console_clear()), - )?; - m.add( - py, - "stdio_thread_get_destination", - cpython::py_fn!(py, stdio_thread_get_destination()), - )?; - m.add( - py, - "stdio_thread_set_destination", - cpython::py_fn!(py, stdio_thread_set_destination(a: PyStdioDestination)), - )?; - - m.add(py, "flush_log", cpython::py_fn!(py, flush_log()))?; - m.add( - py, - "write_log", - cpython::py_fn!(py, write_log(msg: String, level: u64, target: String)), - )?; - m.add( - py, - "set_per_run_log_path", - cpython::py_fn!(py, set_per_run_log_path(a: Option)), - )?; - - m.add( - py, - "task_side_effected", - cpython::py_fn!(py, task_side_effected()), - )?; - m.add( - py, - "teardown_dynamic_ui", - cpython::py_fn!(py, teardown_dynamic_ui(a: PyScheduler, b: PySession)), - )?; - - m.add( - py, - "maybe_set_panic_handler", - cpython::py_fn!(py, maybe_set_panic_handler()), - )?; - m.add( py, "write_digest", @@ -322,59 +273,6 @@ cpython::py_module_initializer!(native_engine, |py, m| { cpython::py_fn!(py, session_isolated_shallow_clone(a: PySession, b: String)), )?; - m.add( - py, - "tasks_task_begin", - cpython::py_fn!( - py, - tasks_task_begin( - tasks: PyTasks, - func: cpython::PyObject, - return_type: cpython::PyType, - side_effecting: bool, - engine_aware_return_type: bool, - cacheable: bool, - name: String, - desc: String, - level: u64 - ) - ), - )?; - m.add( - py, - "tasks_task_end", - cpython::py_fn!(py, tasks_task_end(a: PyTasks)), - )?; - m.add( - py, - "tasks_add_get", - cpython::py_fn!( - py, - tasks_add_get(a: PyTasks, b: cpython::PyType, c: cpython::PyType) - ), - )?; - m.add( - py, - "tasks_add_union", - cpython::py_fn!( - py, - tasks_add_union(a: PyTasks, b: cpython::PyType, c: Vec) - ), - )?; - m.add( - py, - "tasks_add_select", - cpython::py_fn!(py, tasks_add_select(a: PyTasks, b: cpython::PyType)), - )?; - m.add( - py, - "tasks_add_query", - cpython::py_fn!( - py, - tasks_add_query(a: PyTasks, b: cpython::PyType, c: Vec) - ), - )?; - m.add( py, "scheduler_execute", @@ -434,15 +332,6 @@ cpython::py_module_initializer!(native_engine, |py, m| { ), )?; - m.add( - py, - "strongly_connected_components", - cpython::py_fn!( - py, - strongly_connected_components(a: Vec<(cpython::PyObject, Vec)>) - ), - )?; - Ok(()) }); @@ -754,19 +643,16 @@ struct PyResult { engine_traceback: Py, } -fn py_result_from_root(py: Python, result: Result) -> PyO3Result> { +fn py_result_from_root(py: Python, result: Result) -> PyResult { match result { Ok(val) => { let engine_traceback: Vec = vec![]; - Py::new( - py, - PyResult { - is_throw: false, - result: val.into(), - python_traceback: "".into(), - engine_traceback: engine_traceback.into(), - }, - ) + PyResult { + is_throw: false, + result: val.into(), + python_traceback: "".into(), + engine_traceback: engine_traceback.into(), + } } Err(f) => { let (val, python_traceback, engine_traceback) = match f { @@ -784,15 +670,12 @@ fn py_result_from_root(py: Python, result: Result) -> PyO3Result engine_traceback, } => (val, python_traceback, engine_traceback), }; - Py::new( - py, - PyResult { - is_throw: true, - result: val.into(), - python_traceback: python_traceback.into(), - engine_traceback: engine_traceback.into(), - }, - ) + PyResult { + is_throw: true, + result: val.into(), + python_traceback: python_traceback.into(), + engine_traceback: engine_traceback.into(), + } } } } @@ -887,20 +770,21 @@ fn nailgun_server_await_shutdown( }) } +#[pyfunction] fn strongly_connected_components( py: cpython::Python, - adjacency_lists: Vec<(cpython::PyObject, Vec)>, -) -> CPythonPyResult>> { + adjacency_lists: Vec<(PyObject, Vec)>, +) -> PyO3Result>> { let mut graph: DiGraph = Graph::new(); let mut node_ids: HashMap = HashMap::new(); for (node, adjacency_list) in adjacency_lists { - let node_key = Key::from_value(node.clone_ref(py).into())?; + let node_key = Key::from_value(node.into())?; let node_id = *node_ids .entry(node_key) .or_insert_with(|| graph.add_node(node_key)); for dependency in adjacency_list { - let dependency_key = Key::from_value(dependency.clone_ref(py).into())?; + let dependency_key = Key::from_value(dependency.into())?; let dependency_id = node_ids .entry(dependency_key) .or_insert_with(|| graph.add_node(dependency_key)); @@ -1165,14 +1049,14 @@ async fn workunits_to_py_tuple_value<'a>( workunits: impl Iterator, core: &Arc, session: &Session, -) -> CPythonPyResult { +) -> PyO3Result { let mut workunit_values = Vec::new(); for workunit in workunits { let py_value = workunit_to_py_value(workunit, core, session).await?; workunit_values.push(py_value); } - let gil = cpython::Python::acquire_gil(); + let gil = Python::acquire_gil(); Ok(externs::store_tuple(gil.python(), workunit_values)) } @@ -1284,15 +1168,16 @@ fn scheduler_shutdown( Ok(None) } +#[pyfunction] fn scheduler_execute( - py: cpython::Python, - scheduler_ptr: PyScheduler, - session_ptr: PySession, - execution_request_ptr: PyExecutionRequest, -) -> CPythonPyResult { - with_scheduler(py, scheduler_ptr, |scheduler| { - with_execution_request(py, execution_request_ptr, |execution_request| { - with_session(py, session_ptr, |session| { + scheduler_ptr: &PyScheduler, + session_ptr: &PySession, + execution_request_ptr: &PyExecutionRequest, + py: Python, +) -> PyO3Result { + with_scheduler(scheduler_ptr, |scheduler| { + with_execution_request(execution_request_ptr, |execution_request| { + with_session(session_ptr, |session| { // TODO: A parent_id should be an explicit argument. session.workunit_store().init_thread_state(None); py.allow_threads(|| scheduler.execute(execution_request, session)) @@ -1301,7 +1186,7 @@ fn scheduler_execute( .into_iter() .map(|rr| py_result_from_root(py, rr).unwrap().into_object()) .collect::>(); - cpython::PyTuple::new(py, &py_results) + PyTuple::new(py, &py_results) }) .map_err(|e| match e { ExecutionTermination::KeyboardInterrupt => { @@ -1319,46 +1204,44 @@ fn scheduler_execute( }) } +#[pyfunction] fn execution_add_root_select( py: cpython::Python, scheduler_ptr: PyScheduler, - execution_request_ptr: PyExecutionRequest, - param_vals: Vec, - product: cpython::PyType, -) -> PyUnitResult { + execution_request_ptr: &PyExecutionRequest, + param_vals: Vec, + product: &PyType, +) -> PyO3Result<()> { with_scheduler(py, scheduler_ptr, |scheduler| { - with_execution_request(py, execution_request_ptr, |execution_request| { - let product = TypeId::new(&product); + with_execution_request(execution_request_ptr, |execution_request| { + let product = TypeId::new(product); let keys = param_vals .into_iter() .map(|p| Key::from_value(p.into())) .collect::, _>>()?; Params::new(keys) .and_then(|params| scheduler.add_root_select(execution_request, params, product)) - .map_err(|e| cpython::PyErr::new::(py, (e,))) - .map(|()| None) + .map_err(PyException::new_err) }) }) } +#[pyfunction] fn tasks_task_begin( - py: cpython::Python, - tasks_ptr: PyTasks, - func: cpython::PyObject, - output_type: cpython::PyType, + tasks_ptr: &PyTasks, + func: PyObject, + output_type: &PyType, side_effecting: bool, engine_aware_return_type: bool, cacheable: bool, name: String, desc: String, level: u64, -) -> PyUnitResult { - let py_level: PythonLogLevel = level - .try_into() - .map_err(|e| cpython::PyErr::new::(py, (format!("{}", e),)))?; - with_tasks(py, tasks_ptr, |tasks| { +) { + let py_level: PythonLogLevel = level.try_into().map_err(PyException::new_err)?; + with_tasks(tasks_ptr, |tasks| { let func = Function(Key::from_value(func.into())?); - let output_type = TypeId::new(&output_type); + let output_type = TypeId::new(output_type); tasks.task_begin( func, output_type, @@ -1369,76 +1252,56 @@ fn tasks_task_begin( if desc.is_empty() { None } else { Some(desc) }, py_level.into(), ); - Ok(None) }) } -fn tasks_task_end(py: cpython::Python, tasks_ptr: PyTasks) -> PyUnitResult { - with_tasks(py, tasks_ptr, |tasks| { +#[pyfunction] +fn tasks_task_end(tasks_ptr: &PyTasks) { + with_tasks(tasks_ptr, |tasks| { tasks.task_end(); - Ok(None) }) } -fn tasks_add_get( - py: cpython::Python, - tasks_ptr: PyTasks, - output: cpython::PyType, - input: cpython::PyType, -) -> PyUnitResult { - with_tasks(py, tasks_ptr, |tasks| { - let output = TypeId::new(&output); - let input = TypeId::new(&input); +#[pyfunction] +fn tasks_add_get(tasks_ptr: &PyTasks, output: &PyType, input: &PyType) { + with_tasks(tasks_ptr, |tasks| { + let output = TypeId::new(output); + let input = TypeId::new(input); tasks.add_get(output, input); - Ok(None) }) } -fn tasks_add_union( - py: cpython::Python, - tasks_ptr: PyTasks, - output_type: cpython::PyType, - input_types: Vec, -) -> PyUnitResult { - with_tasks(py, tasks_ptr, |tasks| { +#[pyfunction] +fn tasks_add_union(tasks_ptr: &PyTasks, output_type: &PyType, input_types: Vec<&PyType>) { + with_tasks(tasks_ptr, |tasks| { tasks.add_union( - TypeId::new(&output_type), + TypeId::new(output_type), input_types .into_iter() - .map(|type_id| TypeId::new(&type_id)) + .map(|type_id| TypeId::new(type_id)) .collect(), ); - Ok(None) }) } -fn tasks_add_select( - py: cpython::Python, - tasks_ptr: PyTasks, - selector: cpython::PyType, -) -> PyUnitResult { - with_tasks(py, tasks_ptr, |tasks| { - let selector = TypeId::new(&selector); +#[pyfunction] +fn tasks_add_select(tasks_ptr: &PyTasks, selector: &PyType) { + with_tasks(tasks_ptr, |tasks| { + let selector = TypeId::new(selector); tasks.add_select(selector); - Ok(None) }) } -fn tasks_add_query( - py: cpython::Python, - tasks_ptr: PyTasks, - output_type: cpython::PyType, - input_types: Vec, -) -> PyUnitResult { - with_tasks(py, tasks_ptr, |tasks| { +#[pyfunction] +fn tasks_add_query(tasks_ptr: &PyTasks, output_type: &PyType, input_types: Vec<&PyType>) { + with_tasks(tasks_ptr, |tasks| { tasks.query_add( - TypeId::new(&output_type), + TypeId::new(output_type), input_types .into_iter() - .map(|type_id| TypeId::new(&type_id)) + .map(|type_id| TypeId::new(type_id)) .collect(), ); - Ok(None) }) } @@ -1685,9 +1548,10 @@ pub(crate) fn generate_panic_string(payload: &(dyn Any + Send)) -> String { } /// Set up a panic handler, unless RUST_BACKTRACE is set. -fn maybe_set_panic_handler(_: cpython::Python) -> PyUnitResult { +#[pyfunction] +fn maybe_set_panic_handler() { if std::env::var("RUST_BACKTRACE").unwrap_or_else(|_| "0".to_owned()) != "0" { - return Ok(None); + return; } panic::set_hook(Box::new(|panic_info| { let payload = panic_info.payload(); @@ -1703,7 +1567,6 @@ fn maybe_set_panic_handler(_: cpython::Python) -> PyUnitResult { let panic_file_bug_str = "Please set RUST_BACKTRACE=1, re-run, and then file a bug at https://github.com/pantsbuild/pants/issues."; error!("{}", panic_file_bug_str); })); - Ok(None) } fn garbage_collect_store( @@ -1915,31 +1778,22 @@ fn write_digest( }) } +#[pyfunction] fn stdio_initialize( - py: cpython::Python, + py: Python, level: u64, show_rust_3rdparty_logs: bool, show_target: bool, - log_levels_by_target: cpython::PyDict, + log_levels_by_target: HashMap, literal_filters: Vec, regex_filters: Vec, log_file: String, -) -> CPythonPyResult { - let log_levels_by_target = log_levels_by_target - .items(py) - .iter() - .map(|(k, v)| { - let k: String = k.extract(py).unwrap(); - let v: u64 = v.extract(py).unwrap(); - (k, v) - }) - .collect::>(); +) -> PyO3Result { let regex_filters = regex_filters .iter() .map(|re| { Regex::new(re).map_err(|e| { - cpython::PyErr::new::( - py, + PyException::new_err( format!( "Failed to parse warning filter. Please check the global option `--ignore-warnings`.\n\n{}", e, @@ -1958,12 +1812,7 @@ fn stdio_initialize( regex_filters, PathBuf::from(log_file), ) - .map_err(|s| { - cpython::PyErr::new::( - py, - (format!("Could not initialize logging: {}", s),), - ) - })?; + .map_err(|s| PyException::new_err(format!("Could not initialize logging: {}", s)))?; Ok(PyTuple::new( py, @@ -1975,80 +1824,68 @@ fn stdio_initialize( )) } -fn stdio_thread_console_set( - _: cpython::Python, - stdin_fileno: i32, - stdout_fileno: i32, - stderr_fileno: i32, -) -> PyUnitResult { +#[pyfunction] +fn stdio_thread_console_set(stdin_fileno: i32, stdout_fileno: i32, stderr_fileno: i32) { let destination = stdio::new_console_destination(stdin_fileno, stdout_fileno, stderr_fileno); stdio::set_thread_destination(destination); - Ok(None) } -fn stdio_thread_console_color_mode_set(_: cpython::Python, use_color: bool) -> PyUnitResult { +#[pyfunction] +fn stdio_thread_console_color_mode_set(use_color: bool) { stdio::get_destination().stderr_set_use_color(use_color); - Ok(None) } -fn stdio_thread_console_clear(_: cpython::Python) -> PyUnitResult { +#[pyfunction] +fn stdio_thread_console_clear() { stdio::get_destination().console_clear(); - Ok(None) } -fn stdio_thread_get_destination(py: cpython::Python) -> CPythonPyResult { +#[pyfunction] +fn stdio_thread_get_destination(py: Python) -> PyStdioDestination { let dest = stdio::get_destination(); - PyStdioDestination::create_instance(py, dest) + PyStdioDestination(dest) } -fn stdio_thread_set_destination( - py: cpython::Python, - stdio_destination: PyStdioDestination, -) -> PyUnitResult { - stdio::set_thread_destination(stdio_destination.destination(py).clone()); - Ok(None) +#[pyfunction] +fn stdio_thread_set_destination(stdio_destination: PyStdioDestination) { + stdio::set_thread_destination(stdio_destination.0.clone()); } // TODO: Needs to be thread-local / associated with the Console. -fn set_per_run_log_path(py: cpython::Python, log_path: Option) -> PyUnitResult { +#[pyfunction] +fn set_per_run_log_path(py: Python, log_path: Option) { py.allow_threads(|| { PANTS_LOGGER.set_per_run_logs(log_path.map(PathBuf::from)); - Ok(None) }) } -fn write_log(py: cpython::Python, msg: String, level: u64, target: String) -> PyUnitResult { +#[pyfunction] +fn write_log(py: Python, msg: String, level: u64, target: String) { py.allow_threads(|| { Logger::log_from_python(&msg, level, &target).expect("Error logging message"); - Ok(None) }) } -fn task_side_effected(py: cpython::Python) -> PyUnitResult { - nodes::task_side_effected() - .map(|()| None) - .map_err(|e| cpython::PyErr::new::(py, (e,))) +#[pyfunction] +fn task_side_effected() -> PyO3Result<()> { + nodes::task_side_effected().map_err(PyException::new_err) } -fn teardown_dynamic_ui( - py: cpython::Python, - scheduler_ptr: PyScheduler, - session_ptr: PySession, -) -> PyUnitResult { +#[pyfunction] +fn teardown_dynamic_ui(py: cpython::Python, scheduler_ptr: PyScheduler, session_ptr: PySession) { with_scheduler(py, scheduler_ptr, |_scheduler| { with_session(py, session_ptr, |session| { let _ = block_in_place_and_wait(py, || { session.maybe_display_teardown().unit_error().boxed_local() }); - Ok(None) }) }) } -fn flush_log(py: cpython::Python) -> PyUnitResult { +#[pyfunction] +fn flush_log(py: Python) { py.allow_threads(|| { PANTS_LOGGER.flush(); - Ok(None) }) } @@ -2097,9 +1934,7 @@ where scheduler.core.executor.enter(|| f(scheduler)) } -/// /// See `with_scheduler`. -/// fn with_executor(py: cpython::Python, executor_ptr: &PyExecutor, f: F) -> T where F: FnOnce(&Executor) -> T, @@ -2108,9 +1943,7 @@ where f(executor) } -/// /// See `with_scheduler`. -/// fn with_session(py: cpython::Python, session_ptr: PySession, f: F) -> T where F: FnOnce(&Session) -> T, @@ -2119,9 +1952,7 @@ where f(session) } -/// /// See `with_scheduler`. -/// fn with_nailgun_server(py: cpython::Python, nailgun_server_ptr: PyNailgunServer, f: F) -> T where F: FnOnce(&RefCell>, &Executor) -> T, @@ -2131,28 +1962,20 @@ where f(nailgun_server, executor) } -/// /// See `with_scheduler`. -/// -fn with_execution_request( - py: cpython::Python, - execution_request_ptr: PyExecutionRequest, - f: F, -) -> T +fn with_execution_request(execution_request_ptr: &PyExecutionRequest, f: F) -> T where F: FnOnce(&mut ExecutionRequest) -> T, { - let mut execution_request = execution_request_ptr.execution_request(py).borrow_mut(); + let mut execution_request = execution_request_ptr.0.borrow_mut(); f(&mut execution_request) } -/// /// See `with_scheduler`. -/// -fn with_tasks(py: cpython::Python, tasks_ptr: PyTasks, f: F) -> T +fn with_tasks(tasks_ptr: &PyTasks, f: F) -> T where F: FnOnce(&mut Tasks) -> T, { - let mut tasks = tasks_ptr.tasks(py).borrow_mut(); + let mut tasks = tasks_ptr.0.borrow_mut(); f(&mut tasks) }