From 2ac619e5a07a34c91817e8fb3c08a463b5213fad Mon Sep 17 00:00:00 2001
From: Tom Dyas <tom.dyas@gmail.com>
Date: Tue, 1 Oct 2024 21:34:30 -0400
Subject: [PATCH] pyo3: migrate to Bound smart ptr for
 src/externs/engine_aware.rs

---
 src/rust/engine/src/externs/engine_aware.rs | 18 +++++++++---------
 src/rust/engine/src/nodes/mod.rs            |  8 ++++----
 src/rust/engine/src/nodes/task.rs           |  2 +-
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/rust/engine/src/externs/engine_aware.rs b/src/rust/engine/src/externs/engine_aware.rs
index 892e64918ac..7875e4ee55b 100644
--- a/src/rust/engine/src/externs/engine_aware.rs
+++ b/src/rust/engine/src/externs/engine_aware.rs
@@ -21,7 +21,7 @@ use workunit_store::{ArtifactOutput, Level, RunningWorkunit, UserMetadataItem, W
 pub(crate) struct EngineAwareReturnType;
 
 impl EngineAwareReturnType {
-    pub(crate) fn update_workunit(workunit: &mut RunningWorkunit, task_result: &PyAny) {
+    pub(crate) fn update_workunit(workunit: &mut RunningWorkunit, task_result: &Bound<'_, PyAny>) {
         workunit.update_metadata(|old| {
             let new_level = Self::level(task_result);
 
@@ -45,15 +45,15 @@ impl EngineAwareReturnType {
         });
     }
 
-    fn level(obj: &PyAny) -> Option<Level> {
+    fn level(obj: &Bound<'_, PyAny>) -> Option<Level> {
         let level_val = obj.call_method0("level").ok()?;
         if level_val.is_none() {
             return None;
         }
-        externs::val_to_log_level(level_val).ok()
+        externs::val_to_log_level(level_val.as_gil_ref()).ok()
     }
 
-    fn message(obj: &PyAny) -> Option<String> {
+    fn message(obj: &Bound<'_, PyAny>) -> Option<String> {
         let msg_val = obj.call_method0("message").ok()?;
         if msg_val.is_none() {
             return None;
@@ -61,7 +61,7 @@ impl EngineAwareReturnType {
         msg_val.extract().ok()
     }
 
-    fn artifacts(obj: &PyAny) -> Option<Vec<(String, ArtifactOutput)>> {
+    fn artifacts(obj: &Bound<'_, PyAny>) -> Option<Vec<(String, ArtifactOutput)>> {
         let artifacts_val = obj.call_method0("artifacts").ok()?;
         if artifacts_val.is_none() {
             return None;
@@ -84,7 +84,7 @@ impl EngineAwareReturnType {
         Some(output)
     }
 
-    pub(crate) fn is_cacheable(obj: &PyAny) -> Option<bool> {
+    pub(crate) fn is_cacheable(obj: &Bound<'_, PyAny>) -> Option<bool> {
         obj.call_method0("cacheable").ok()?.extract().ok()
     }
 }
@@ -92,7 +92,7 @@ impl EngineAwareReturnType {
 pub struct EngineAwareParameter;
 
 impl EngineAwareParameter {
-    pub fn debug_hint(obj: &PyAny) -> Option<String> {
+    pub fn debug_hint(obj: &Bound<'_, PyAny>) -> Option<String> {
         let hint = obj.call_method0("debug_hint").ok()?;
         if hint.is_none() {
             return None;
@@ -100,12 +100,12 @@ impl EngineAwareParameter {
         hint.extract().ok()
     }
 
-    pub fn metadata(obj: &PyAny) -> Vec<(String, UserMetadataItem)> {
+    pub fn metadata(obj: &Bound<'_, PyAny>) -> Vec<(String, UserMetadataItem)> {
         metadata_for(obj).unwrap_or_default()
     }
 }
 
-fn metadata_for(obj: &PyAny) -> Option<Vec<(String, UserMetadataItem)>> {
+fn metadata_for(obj: &Bound<'_, PyAny>) -> Option<Vec<(String, UserMetadataItem)>> {
     let metadata_val = obj.call_method0("metadata").ok()?;
     if metadata_val.is_none() {
         return None;
diff --git a/src/rust/engine/src/nodes/mod.rs b/src/rust/engine/src/nodes/mod.rs
index bbce069dc13..eccba4bc826 100644
--- a/src/rust/engine/src/nodes/mod.rs
+++ b/src/rust/engine/src/nodes/mod.rs
@@ -396,7 +396,7 @@ impl NodeKey {
 
                 let displayable_param_names: Vec<_> = Python::with_gil(|py| {
                     Self::engine_aware_params(context, py, &task.params)
-                        .filter_map(|k| EngineAwareParameter::debug_hint((*k.value).as_ref(py)))
+                        .filter_map(|k| EngineAwareParameter::debug_hint((*k.value).bind(py)))
                         .collect()
                 });
 
@@ -502,7 +502,7 @@ impl Node for NodeKey {
                 if let Some(params) = maybe_params {
                     Python::with_gil(|py| {
                         Self::engine_aware_params(&context, py, params)
-                            .flat_map(|k| EngineAwareParameter::metadata((*k.value).as_ref(py)))
+                            .flat_map(|k| EngineAwareParameter::metadata((*k.value).bind(py)))
                             .collect()
                     })
                 } else {
@@ -601,7 +601,7 @@ impl Node for NodeKey {
             }
             (NodeKey::Task(ref t), NodeOutput::Value(ref v)) if t.task.engine_aware_return_type => {
                 Python::with_gil(|py| {
-                    EngineAwareReturnType::is_cacheable((**v).as_ref(py)).unwrap_or(true)
+                    EngineAwareReturnType::is_cacheable((**v).bind(py)).unwrap_or(true)
                 })
             }
             _ => true,
@@ -652,7 +652,7 @@ impl Display for NodeKey {
                             .keys()
                             .filter_map(|k| {
                                 EngineAwareParameter::debug_hint(
-                                    k.to_value().clone_ref(py).into_ref(py),
+                                    k.to_value().clone_ref(py).bind(py),
                                 )
                             })
                             .collect::<Vec<_>>()
diff --git a/src/rust/engine/src/nodes/task.rs b/src/rust/engine/src/nodes/task.rs
index 588f96446dc..9109f20feb5 100644
--- a/src/rust/engine/src/nodes/task.rs
+++ b/src/rust/engine/src/nodes/task.rs
@@ -331,7 +331,7 @@ impl Task {
 
         if self.task.engine_aware_return_type {
             Python::with_gil(|py| {
-                EngineAwareReturnType::update_workunit(workunit, (*result_val).as_ref(py))
+                EngineAwareReturnType::update_workunit(workunit, (*result_val).bind(py))
             })
         };