From 73e9a3067a15ee85a1548d65d9888e8ddb8fe2f9 Mon Sep 17 00:00:00 2001
From: Itamar Turner-Trauring <itamar@itamarst.org>
Date: Tue, 21 May 2024 13:53:56 -0400
Subject: [PATCH] Switch to newer PyO3 0.21 APIs (Bound<> and friends) (#82)

Co-authored-by: Itamar Turner-Trauring <itamar@pythonspeed.com>
---
 .../extend_polars/src/lib.rs                  |  2 +-
 pyo3-polars/src/ffi/to_rust.rs                |  2 +-
 pyo3-polars/src/lib.rs                        | 19 ++++++++++---------
 rust-toolchain.toml                           |  2 +-
 4 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/example/extend_polars_python_dispatch/extend_polars/src/lib.rs b/example/extend_polars_python_dispatch/extend_polars/src/lib.rs
index 1bc0ff6..fac8510 100644
--- a/example/extend_polars_python_dispatch/extend_polars/src/lib.rs
+++ b/example/extend_polars_python_dispatch/extend_polars/src/lib.rs
@@ -31,7 +31,7 @@ fn lazy_parallel_jaccard(pydf: PyLazyFrame, col_a: &str, col_b: &str) -> PyResul
 
 /// A Python module implemented in Rust.
 #[pymodule]
-fn extend_polars(_py: Python, m: &PyModule) -> PyResult<()> {
+fn extend_polars(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(parallel_jaccard, m)?)?;
     m.add_function(wrap_pyfunction!(lazy_parallel_jaccard, m)?)?;
     m.add_function(wrap_pyfunction!(debug, m)?)?;
diff --git a/pyo3-polars/src/ffi/to_rust.rs b/pyo3-polars/src/ffi/to_rust.rs
index 27f5037..0590d09 100644
--- a/pyo3-polars/src/ffi/to_rust.rs
+++ b/pyo3-polars/src/ffi/to_rust.rs
@@ -4,7 +4,7 @@ use polars::prelude::*;
 use pyo3::ffi::Py_uintptr_t;
 use pyo3::prelude::*;
 
-pub fn array_to_rust(obj: &PyAny) -> PyResult<ArrayRef> {
+pub fn array_to_rust(obj: &Bound<PyAny>) -> PyResult<ArrayRef> {
     // prepare a pointer to receive the Array struct
     let array = Box::new(ffi::ArrowArray::empty());
     let schema = Box::new(ffi::ArrowSchema::empty());
diff --git a/pyo3-polars/src/lib.rs b/pyo3-polars/src/lib.rs
index eec501f..08aa0c5 100644
--- a/pyo3-polars/src/lib.rs
+++ b/pyo3-polars/src/lib.rs
@@ -23,7 +23,7 @@
 //!
 //! /// A Python module implemented in Rust.
 //! #[pymodule]
-//! fn expression_lib(_py: Python, m: &PyModule) -> PyResult<()> {
+//! fn expression_lib(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
 //!     m.add_function(wrap_pyfunction!(my_cool_function, m)?)?;
 //!     Ok(())
 //! }
@@ -119,22 +119,23 @@ impl AsRef<LazyFrame> for PyLazyFrame {
 }
 
 impl<'a> FromPyObject<'a> for PySeries {
-    fn extract(ob: &'a PyAny) -> PyResult<Self> {
+    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
         let ob = ob.call_method0("rechunk")?;
 
         let name = ob.getattr("name")?;
-        let name = name.str()?.to_str()?;
+        let py_name = name.str()?;
+        let name = py_name.to_cow()?;
 
         let arr = ob.call_method0("to_arrow")?;
-        let arr = ffi::to_rust::array_to_rust(arr)?;
+        let arr = ffi::to_rust::array_to_rust(&arr)?;
         Ok(PySeries(
-            Series::try_from((name, arr)).map_err(PyPolarsErr::from)?,
+            Series::try_from((&*name, arr)).map_err(PyPolarsErr::from)?,
         ))
     }
 }
 
 impl<'a> FromPyObject<'a> for PyDataFrame {
-    fn extract(ob: &'a PyAny) -> PyResult<Self> {
+    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
         let series = ob.call_method0("get_columns")?;
         let n = ob.getattr("width")?.extract::<usize>()?;
         let mut columns = Vec::with_capacity(n);
@@ -149,7 +150,7 @@ impl<'a> FromPyObject<'a> for PyDataFrame {
 
 #[cfg(feature = "lazy")]
 impl<'a> FromPyObject<'a> for PyLazyFrame {
-    fn extract(ob: &'a PyAny) -> PyResult<Self> {
+    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
         let s = ob.call_method0("__getstate__")?.extract::<Vec<u8>>()?;
         let lp: LogicalPlan = ciborium::de::from_reader(&*s).map_err(
             |e| PyPolarsErr::Other(
@@ -237,9 +238,9 @@ impl IntoPy<PyObject> for PyDataFrame {
 #[cfg(feature = "lazy")]
 impl IntoPy<PyObject> for PyLazyFrame {
     fn into_py(self, py: Python<'_>) -> PyObject {
-        let polars = py.import("polars").expect("polars not installed");
+        let polars = py.import_bound("polars").expect("polars not installed");
         let cls = polars.getattr("LazyFrame").unwrap();
-        let instance = cls.call_method1("__new__", (cls,)).unwrap();
+        let instance = cls.call_method1("__new__", (&cls,)).unwrap();
         let mut writer: Vec<u8> = vec![];
         ciborium::ser::into_writer(&self.0.logical_plan, &mut writer).unwrap();
 
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
index 9b6fc43..71f1e88 100644
--- a/rust-toolchain.toml
+++ b/rust-toolchain.toml
@@ -1,2 +1,2 @@
 [toolchain]
-channel = "nightly-2024-03-28"
+channel = "nightly-2024-05-14"