Use cpython types in Rust functions that manipulate python objects (#10942)

### Problem

We have a number of Rust functions in the `externs` module that manipulate Python values that make use of the `Value` type, which is our own reference-counted wrapper around the `PyObject` type provided by the `cpython` library. The `Value` type predates the use of `cpython` for interacting with Python code, and `PyObject` is a more flexible type than the previous type(s) used to interact with Python. Now it is often the case that we take `Value`s returned from one of the `externs` functions and immediately turn it back into a `PyObject`. This makes the API more confusing to reason about, and forces unnecessary and potentially costly type conversions.

### Solution

Implement `AsRef<PyObject>` for `Value` so we can cheaply turn existing parameters of type `Value` into `&PyObject`s. Then, switch out as many function signatures as possible to work with `PyObject`s, using the `AsRef` and other existing conversion traits to turn `PyObject`s into `Value`s where necessary, and lets us save type conversions when we don't need them.

Author: gshuflin

src/rust/engine/src/core.rs

@@ -3,6 +3,7 @@
 
 use fnv::FnvHasher;
 
+use std::convert::AsRef;
 use std::ops::Deref;
 use std::sync::Arc;
 use std::{fmt, hash};
@@ -264,6 +265,12 @@ impl Deref for Value {
   }
 }
 
+impl AsRef<PyObject> for Value {
+  fn as_ref(&self) -> &PyObject {
+    &self.0
+  }
+}
+
 impl fmt::Debug for Value {
   fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     write!(f, "{}", externs::val_to_str(&self))
@@ -347,7 +354,12 @@ impl Failure {
 }
 
 impl Failure {
-  pub fn from_py_err(py: Python, mut py_err: PyErr) -> Failure {
+  pub fn from_py_err(py_err: PyErr) -> Failure {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+    Failure::from_py_err_with_gil(py, py_err)
+  }
+  pub fn from_py_err_with_gil(py: Python, mut py_err: PyErr) -> Failure {
     let val = Value::from(py_err.instance(py));
     let python_traceback = if let Some(tb) = py_err.ptraceback.as_ref() {
       let locals = PyDict::new(py);

src/rust/engine/src/externs/engine_aware.rs

@@ -4,6 +4,7 @@
 use crate::core::Value;
 use crate::externs;
 use crate::nodes::lift_directory_digest;
+use crate::Failure;
 use crate::Types;
 
 use cpython::{PyDict, PyObject, PyString, Python};
@@ -24,9 +25,9 @@ impl EngineAwareInformation for EngineAwareLevel {
   type MaybeOutput = Level;
 
   fn retrieve(_types: &Types, value: &Value) -> Option<Level> {
-    let new_level_val: Value = externs::call_method(&value, "level", &[]).ok()?;
+    let new_level_val = externs::call_method(value.as_ref(), "level", &[]).ok()?;
     let new_level_val = externs::check_for_python_none(new_level_val)?;
-    externs::val_to_log_level(&new_level_val).ok()
+    externs::val_to_log_level(&new_level_val.into()).ok()
   }
 }
 
@@ -36,9 +37,9 @@ impl EngineAwareInformation for Message {
   type MaybeOutput = String;
 
   fn retrieve(_types: &Types, value: &Value) -> Option<String> {
-    let msg_val: Value = externs::call_method(&value, "message", &[]).ok()?;
+    let msg_val = externs::call_method(&value, "message", &[]).ok()?;
     let msg_val = externs::check_for_python_none(msg_val)?;
-    Some(externs::val_to_str(&msg_val))
+    Some(externs::val_to_str(&msg_val.into()))
   }
 }
 
@@ -48,14 +49,15 @@ impl EngineAwareInformation for Artifacts {
   type MaybeOutput = Vec<(String, Digest)>;
 
   fn retrieve(types: &Types, value: &Value) -> Option<Self::MaybeOutput> {
-    let artifacts_val: Value = match externs::call_method(&value, "artifacts", &[]) {
+    let artifacts_val = match externs::call_method(&value, "artifacts", &[]) {
       Ok(value) => value,
-      Err(e) => {
-        log::error!("Error calling `artifacts` method: {}", e);
+      Err(py_err) => {
+        let failure = Failure::from_py_err(py_err);
+        log::error!("Error calling `artifacts` method: {}", failure);
         return None;
       }
     };
-    let artifacts_val: Value = externs::check_for_python_none(artifacts_val)?;
+    let artifacts_val = externs::check_for_python_none(artifacts_val)?;
     let gil = Python::acquire_gil();
     let py = gil.python();
     let artifacts_dict: &PyDict = artifacts_val.cast_as::<PyDict>(py).ok()?;
@@ -72,7 +74,7 @@ impl EngineAwareInformation for Artifacts {
           return None;
         }
       };
-      let digest_value: PyObject = externs::getattr(&Value::new(value), "digest")
+      let digest_value: PyObject = externs::getattr(&value, "digest")
         .map_err(|e| {
           log::error!("Error in EngineAware.artifacts() - no `digest` attr: {}", e);
         })
@@ -100,6 +102,6 @@ impl EngineAwareInformation for DebugHint {
     externs::call_method(&value, "debug_hint", &[])
       .ok()
       .and_then(externs::check_for_python_none)
-      .map(|val| externs::val_to_str(&val))
+      .map(|val| externs::val_to_str(&val.into()))
   }
 }

src/rust/engine/src/externs/interface.rs

@@ -692,8 +692,9 @@ fn nailgun_server_create(
           stdout_fd,
           stderr_fd,
         ];
-        match externs::call(&runner, &runner_args) {
-          Ok(exit_code_val) => {
+        match externs::call_function(&runner, &runner_args) {
+          Ok(exit_code) => {
+            let exit_code_val: Value = exit_code.into();
             // TODO: We don't currently expose a "project_i32", but it will not be necessary with
             // https://github.com/pantsbuild/pants/pull/9593.
             nailgun::ExitCode(externs::val_to_str(&exit_code_val).parse().unwrap())

src/rust/engine/src/externs/mod.rs

@@ -15,6 +15,7 @@ mod interface;
 mod interface_tests;
 
 use std::collections::BTreeMap;
+use std::convert::AsRef;
 use std::convert::TryInto;
 use std::fmt;
 use std::sync::atomic;
@@ -97,14 +98,6 @@ pub fn type_for(py_type: PyType) -> TypeId {
   })
 }
 
-pub fn acquire_key_for(val: Value) -> Result<Key, Failure> {
-  key_for(val).map_err(|e| {
-    let gil = Python::acquire_gil();
-    let py = gil.python();
-    Failure::from_py_err(py, e)
-  })
-}
-
 pub fn key_for(val: Value) -> Result<Key, PyErr> {
   with_interns_mut(|interns| {
     let gil = Python::acquire_gil();
@@ -174,7 +167,7 @@ pub fn store_bool(val: bool) -> Value {
 ///
 /// Check if a Python object has the specified field.
 ///
-pub fn hasattr(value: &Value, field: &str) -> bool {
+pub fn hasattr(value: &PyObject, field: &str) -> bool {
   let gil = Python::acquire_gil();
   let py = gil.python();
   value.hasattr(py, field).unwrap()
@@ -183,7 +176,7 @@ pub fn hasattr(value: &Value, field: &str) -> bool {
 ///
 /// Gets an attribute of the given value as the given type.
 ///
-fn getattr<T>(value: &Value, field: &str) -> Result<T, String>
+fn getattr<T>(value: &PyObject, field: &str) -> Result<T, String>
 where
   for<'a> T: FromPyObject<'a>,
 {
@@ -235,27 +228,27 @@ fn collect_iterable(value: &Value) -> Result<Vec<Value>, String> {
 /// Pulls out the value specified by the field name from a given Value
 ///
 pub fn project_ignoring_type(value: &Value, field: &str) -> Value {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_iterable(value: &Value) -> Vec<Value> {
   collect_iterable(value).unwrap()
 }
 
 pub fn project_multi(value: &Value, field: &str) -> Vec<Value> {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_bool(value: &Value, field: &str) -> bool {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_multi_strs(value: &Value, field: &str) -> Vec<String> {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_frozendict(value: &Value, field: &str) -> BTreeMap<String, String> {
-  let frozendict = Value::new(getattr(value, field).unwrap());
+  let frozendict = getattr(value.as_ref(), field).unwrap();
   let pydict: PyDict = getattr(&frozendict, "_data").unwrap();
   let gil = Python::acquire_gil();
   let py = gil.python();
@@ -271,25 +264,25 @@ pub fn project_str(value: &Value, field: &str) -> String {
   // cloning in some cases.
   // TODO: We can't directly extract as a string here, because val_to_str defaults to empty string
   // for None.
-  val_to_str(&getattr(value, field).unwrap())
+  val_to_str(&getattr(value.as_ref(), field).unwrap())
 }
 
 pub fn project_u64(value: &Value, field: &str) -> u64 {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_maybe_u64(value: &Value, field: &str) -> Result<u64, String> {
-  getattr(value, field)
+  getattr(value.as_ref(), field)
 }
 
 pub fn project_f64(value: &Value, field: &str) -> f64 {
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn project_bytes(value: &Value, field: &str) -> Vec<u8> {
   // TODO: It's possible to view a python bytes as a `&[u8]`, so we could avoid actually
   // cloning in some cases.
-  getattr(value, field).unwrap()
+  getattr(value.as_ref(), field).unwrap()
 }
 
 pub fn key_to_str(key: &Key) -> String {
@@ -327,41 +320,31 @@ pub fn create_exception(msg: &str) -> Value {
   Value::from(with_externs(|py, e| e.call_method(py, "create_exception", (msg,), None)).unwrap())
 }
 
-pub fn check_for_python_none(value: Value) -> Option<Value> {
+pub fn check_for_python_none(value: PyObject) -> Option<PyObject> {
   let gil = Python::acquire_gil();
   let py = gil.python();
 
-  if *value == py.None() {
+  if value == py.None() {
     return None;
   }
   Some(value)
 }
 
-pub fn call_method(value: &Value, method: &str, args: &[Value]) -> Result<Value, Failure> {
+pub fn call_method(value: &PyObject, method: &str, args: &[Value]) -> Result<PyObject, PyErr> {
   let arg_handles: Vec<PyObject> = args.iter().map(|v| v.clone().into()).collect();
   let gil = Python::acquire_gil();
   let args_tuple = PyTuple::new(gil.python(), &arg_handles);
-  value
-    .call_method(gil.python(), method, args_tuple, None)
-    .map(Value::from)
-    .map_err(|py_err| Failure::from_py_err(gil.python(), py_err))
+  value.call_method(gil.python(), method, args_tuple, None)
 }
 
-pub fn call_function(func: &Value, args: &[Value]) -> Result<PyObject, PyErr> {
+pub fn call_function<T: AsRef<PyObject>>(func: T, args: &[Value]) -> Result<PyObject, PyErr> {
+  let func: &PyObject = func.as_ref();
   let arg_handles: Vec<PyObject> = args.iter().map(|v| v.clone().into()).collect();
   let gil = Python::acquire_gil();
   let args_tuple = PyTuple::new(gil.python(), &arg_handles);
   func.call(gil.python(), args_tuple, None)
 }
 
-pub fn call(func: &Value, args: &[Value]) -> Result<Value, Failure> {
-  let output = call_function(func, args);
-  output.map(Value::from).map_err(|py_err| {
-    let gil = Python::acquire_gil();
-    Failure::from_py_err(gil.python(), py_err)
-  })
-}
-
 pub fn generator_send(generator: &Value, arg: &Value) -> Result<GeneratorResponse, Failure> {
   let response = with_externs(|py, e| {
     e.call_method(
@@ -370,7 +353,7 @@ pub fn generator_send(generator: &Value, arg: &Value) -> Result<GeneratorRespons
       (generator as &PyObject, arg as &PyObject),
       None,
     )
-    .map_err(|py_err| Failure::from_py_err(py, py_err))
+    .map_err(|py_err| Failure::from_py_err_with_gil(py, py_err))
   })?;
 
   let gil = Python::acquire_gil();
@@ -398,7 +381,7 @@ pub fn generator_send(generator: &Value, arg: &Value) -> Result<GeneratorRespons
         .map(|g| {
           let get = g
             .cast_as::<PyGeneratorResponseGet>(py)
-            .map_err(|e| Failure::from_py_err(py, e.into()))?;
+            .map_err(|e| Failure::from_py_err_with_gil(py, e.into()))?;
           Ok(Get::new(py, interns, get)?)
         })
         .collect::<Result<Vec<_>, _>>()?;
@@ -520,7 +503,7 @@ impl Get {
       output: interns.type_insert(py, get.product(py).clone_ref(py)),
       input: interns
         .key_insert(py, get.subject(py).clone_ref(py).into())
-        .map_err(|e| Failure::from_py_err(py, e))?,
+        .map_err(|e| Failure::from_py_err_with_gil(py, e))?,
       input_type: Some(interns.type_insert(py, get.declared_subject(py).clone_ref(py))),
     })
   }

src/rust/engine/src/intrinsics.rs

@@ -7,6 +7,7 @@ use crate::nodes::{
 };
 use crate::tasks::Intrinsic;
 use crate::types::Types;
+use crate::Failure;
 
 use fs::RelativePath;
 use futures::compat::Future01CompatExt;
@@ -304,7 +305,7 @@ fn download_file_to_digest(
   mut args: Vec<Value>,
 ) -> BoxFuture<'static, NodeResult<Value>> {
   async move {
-    let key = externs::acquire_key_for(args.pop().unwrap())?;
+    let key = externs::key_for(args.pop().unwrap()).map_err(Failure::from_py_err)?;
     let digest = context.get(DownloadedFile(key)).await?;
     Snapshot::store_directory_digest(&digest).map_err(|s| throw(&s))
   }
@@ -354,7 +355,7 @@ fn create_digest_to_digest(
         let path = RelativePath::new(PathBuf::from(path))
           .map_err(|e| format!("The `path` must be relative: {:?}", e))?;
 
-        if externs::hasattr(&file_content_or_directory, "content") {
+        if externs::hasattr(file_content_or_directory.as_ref(), "content") {
           let bytes = bytes::Bytes::from(externs::project_bytes(
             &file_content_or_directory,
             "content",

src/rust/engine/src/nodes.rs

@@ -910,13 +910,13 @@ impl Task {
               .cloned()
               .ok_or_else(|| match get.input_type {
                 Some(ty) if externs::is_union(ty) => {
-                  let value = externs::type_for_type_id(ty).into_object().into();
+                  let value = externs::type_for_type_id(ty).into_object();
                   match externs::call_method(
                     &value,
                     "non_member_error_message",
                     &[externs::val_for(&get.input)],
                   ) {
-                    Ok(err_msg) => throw(&externs::val_to_str(&err_msg)),
+                    Ok(err_msg) => throw(&externs::val_to_str(&err_msg.into())),
                     // If the non_member_error_message() call failed for any reason,
                     // fall back to a generic message.
                     Err(_e) => throw(&format!(
@@ -1025,7 +1025,9 @@ impl WrappedNode for Task {
     let product = self.product;
     let can_modify_workunit = self.task.can_modify_workunit;
 
-    let mut result_val = externs::call(&externs::val_for(&func.0), &deps)?;
+    let result_val =
+      externs::call_function(&externs::val_for(&func.0), &deps).map_err(Failure::from_py_err)?;
+    let mut result_val: Value = result_val.into();
     let mut result_type = externs::get_type_for(&result_val);
     if result_type == context.core.types.coroutine {
       result_val = Self::generate(context.clone(), params, entry, result_val).await?;

src/rust/engine/src/scheduler.rs

@@ -591,7 +591,7 @@ fn maybe_break_execution_loop(python_signal_fn: &Value) -> Option<ExecutionTermi
       {
         Some(ExecutionTermination::KeyboardInterrupt)
       } else {
-        let failure = Failure::from_py_err(py, e);
+        let failure = Failure::from_py_err_with_gil(py, e);
         std::mem::drop(gil);
         Some(ExecutionTermination::Fatal(format!(
           "Error when checking Python signal state: {}",