chore(python): Continue converting to PyO3 0.21 Bound<> APIs (pola-rs…

…#16081) Co-authored-by: Itamar Turner-Trauring <[email protected]>
orlp · May 9, 2024 · 1d13d62 · 1d13d62
1 parent 44852b4
commit 1d13d62
Show file tree

Hide file tree

Showing 14 changed files with 63 additions and 56 deletions.
diff --git a/py-polars/src/arrow_interop/to_rust.rs b/py-polars/src/arrow_interop/to_rust.rs
@@ -9,7 +9,7 @@ use pyo3::types::PyList;
 
 use crate::error::PyPolarsErr;
 
-pub fn field_to_rust(obj: &PyAny) -> PyResult<Field> {
+pub fn field_to_rust(obj: Bound<'_, PyAny>) -> PyResult<Field> {
     let schema = Box::new(ffi::ArrowSchema::empty());
     let schema_ptr = &*schema as *const ffi::ArrowSchema;
 
@@ -20,7 +20,7 @@ pub fn field_to_rust(obj: &PyAny) -> PyResult<Field> {
 }
 
 // PyList<Field> which you get by calling `list(schema)`
-pub fn pyarrow_schema_to_rust(obj: &PyList) -> PyResult<Schema> {
+pub fn pyarrow_schema_to_rust(obj: &Bound<'_, PyList>) -> PyResult<Schema> {
     obj.into_iter().map(field_to_rust).collect()
 }
 

diff --git a/py-polars/src/conversion/mod.rs b/py-polars/src/conversion/mod.rs
@@ -61,12 +61,12 @@ impl<T> From<T> for Wrap<T> {
 }
 
 // extract a Rust DataFrame from a python DataFrame, that is DataFrame<PyDataFrame<RustDataFrame>>
-pub(crate) fn get_df(obj: &PyAny) -> PyResult<DataFrame> {
+pub(crate) fn get_df(obj: &Bound<'_, PyAny>) -> PyResult<DataFrame> {
     let pydf = obj.getattr(intern!(obj.py(), "_df"))?;
     Ok(pydf.extract::<PyDataFrame>()?.df)
 }
 
-pub(crate) fn get_lf(obj: &PyAny) -> PyResult<LazyFrame> {
+pub(crate) fn get_lf(obj: &Bound<'_, PyAny>) -> PyResult<LazyFrame> {
     let pydf = obj.getattr(intern!(obj.py(), "_ldf"))?;
     Ok(pydf.extract::<PyLazyFrame>()?.ldf)
 }
@@ -455,7 +455,7 @@ impl FromPyObject<'_> for Wrap<Schema> {
 
 impl IntoPy<PyObject> for Wrap<&Schema> {
     fn into_py(self, py: Python<'_>) -> PyObject {
-        let dict = PyDict::new(py);
+        let dict = PyDict::new_bound(py);
         for (k, v) in self.0.iter() {
             dict.set_item(k.as_str(), Wrap(v.clone())).unwrap();
         }

diff --git a/py-polars/src/error.rs b/py-polars/src/error.rs
@@ -123,12 +123,12 @@ impl IntoPy<PyObject> for Wrap<PolarsWarning> {
     fn into_py(self, py: Python<'_>) -> PyObject {
         match self.0 {
             PolarsWarning::CategoricalRemappingWarning => {
-                CategoricalRemappingWarning::type_object(py).to_object(py)
+                CategoricalRemappingWarning::type_object_bound(py).to_object(py)
             },
             PolarsWarning::MapWithoutReturnDtypeWarning => {
-                MapWithoutReturnDtypeWarning::type_object(py).to_object(py)
+                MapWithoutReturnDtypeWarning::type_object_bound(py).to_object(py)
             },
-            PolarsWarning::UserWarning => PyUserWarning::type_object(py).to_object(py),
+            PolarsWarning::UserWarning => PyUserWarning::type_object_bound(py).to_object(py),
         }
     }
 }
diff --git a/py-polars/src/expr/array.rs b/py-polars/src/expr/array.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use polars::prelude::*;
 use polars_ops::prelude::array::ArrToStructNameGenerator;
 use pyo3::prelude::*;
@@ -112,7 +114,7 @@ impl PyExpr {
             Arc::new(move |idx: usize| {
                 Python::with_gil(|py| {
                     let out = lambda.call1(py, (idx,)).unwrap();
-                    let out: SmartString = out.extract::<&str>(py).unwrap().into();
+                    let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
                     out
                 })
             }) as ArrToStructNameGenerator

diff --git a/py-polars/src/expr/general.rs b/py-polars/src/expr/general.rs
@@ -1,3 +1,4 @@
+use std::io::Cursor;
 use std::ops::Neg;
 
 use polars::lazy::dsl;
@@ -6,6 +7,7 @@ use polars::series::ops::NullBehavior;
 use polars_core::series::IsSorted;
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
 use pyo3::types::PyBytes;
 
 use crate::conversion::{parse_fill_null_strategy, vec_extract_wrapped, Wrap};
@@ -83,14 +85,15 @@ impl PyExpr {
         ciborium::ser::into_writer(&self.inner, &mut writer)
             .map_err(|e| PyPolarsErr::Other(format!("{}", e)))?;
 
-        Ok(PyBytes::new(py, &writer).to_object(py))
+        Ok(PyBytes::new_bound(py, &writer).to_object(py))
     }
 
     fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
         // Used in pickle/pickling
-        match state.extract::<&PyBytes>(py) {
+        match state.extract::<PyBackedBytes>(py) {
             Ok(s) => {
-                self.inner = ciborium::de::from_reader(s.as_bytes())
+                let cursor = Cursor::new(&*s);
+                self.inner = ciborium::de::from_reader(cursor)
                     .map_err(|e| PyPolarsErr::Other(format!("{}", e)))?;
                 Ok(())
             },

diff --git a/py-polars/src/expr/list.rs b/py-polars/src/expr/list.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use polars::prelude::*;
 use polars::series::ops::NullBehavior;
 use pyo3::prelude::*;
@@ -212,7 +214,7 @@ impl PyExpr {
             Arc::new(move |idx: usize| {
                 Python::with_gil(|py| {
                     let out = lambda.call1(py, (idx,)).unwrap();
-                    let out: SmartString = out.extract::<&str>(py).unwrap().into();
+                    let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
                     out
                 })
             }) as NameGenerator

diff --git a/py-polars/src/expr/name.rs b/py-polars/src/expr/name.rs
@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use polars::prelude::*;
 use pyo3::prelude::*;
 use smartstring::alias::String as SmartString;
@@ -46,7 +48,7 @@ impl PyExpr {
         let name_mapper = Arc::new(move |name: &str| {
             Python::with_gil(|py| {
                 let out = name_mapper.call1(py, (name,)).unwrap();
-                let out: SmartString = out.extract::<&str>(py).unwrap().into();
+                let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
                 out
             })
         }) as FieldsNameMapper;

diff --git a/py-polars/src/expr/rolling.rs b/py-polars/src/expr/rolling.rs
@@ -244,7 +244,7 @@ impl PyExpr {
                     },
                     Err(_) => {
                         let obj = out;
-                        let is_float = obj.as_ref(py).is_instance_of::<PyFloat>();
+                        let is_float = obj.bind(py).is_instance_of::<PyFloat>();
 
                         let dtype = s.dtype();
 

diff --git a/py-polars/src/functions/eager.rs b/py-polars/src/functions/eager.rs
@@ -7,20 +7,20 @@ use crate::error::PyPolarsErr;
 use crate::{PyDataFrame, PySeries};
 
 #[pyfunction]
-pub fn concat_df(dfs: &PyAny, py: Python) -> PyResult<PyDataFrame> {
+pub fn concat_df(dfs: &Bound<'_, PyAny>, py: Python) -> PyResult<PyDataFrame> {
     use polars_core::error::PolarsResult;
     use polars_core::utils::rayon::prelude::*;
 
     let mut iter = dfs.iter()?;
     let first = iter.next().unwrap()?;
 
-    let first_rdf = get_df(first)?;
+    let first_rdf = get_df(&first)?;
     let identity_df = first_rdf.clear();
 
     let mut rdfs: Vec<PolarsResult<DataFrame>> = vec![Ok(first_rdf)];
 
     for item in iter {
-        let rdf = get_df(item?)?;
+        let rdf = get_df(&item?)?;
         rdfs.push(Ok(rdf));
     }
 
@@ -63,13 +63,13 @@ pub fn concat_series(series: &Bound<'_, PyAny>) -> PyResult<PySeries> {
 }
 
 #[pyfunction]
-pub fn concat_df_diagonal(dfs: &PyAny) -> PyResult<PyDataFrame> {
+pub fn concat_df_diagonal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
     let iter = dfs.iter()?;
 
     let dfs = iter
         .map(|item| {
             let item = item?;
-            get_df(item)
+            get_df(&item)
         })
         .collect::<PyResult<Vec<_>>>()?;
 
@@ -78,13 +78,13 @@ pub fn concat_df_diagonal(dfs: &PyAny) -> PyResult<PyDataFrame> {
 }
 
 #[pyfunction]
-pub fn concat_df_horizontal(dfs: &PyAny) -> PyResult<PyDataFrame> {
+pub fn concat_df_horizontal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
     let iter = dfs.iter()?;
 
     let dfs = iter
         .map(|item| {
             let item = item?;
-            get_df(item)
+            get_df(&item)
         })
         .collect::<PyResult<Vec<_>>>()?;
 

diff --git a/py-polars/src/functions/lazy.rs b/py-polars/src/functions/lazy.rs
@@ -153,7 +153,7 @@ pub fn cols(names: Vec<String>) -> PyExpr {
 
 #[pyfunction]
 pub fn concat_lf(
-    seq: &PyAny,
+    seq: &Bound<'_, PyAny>,
     rechunk: bool,
     parallel: bool,
     to_supertypes: bool,
@@ -163,7 +163,7 @@ pub fn concat_lf(
 
     for res in seq.iter()? {
         let item = res?;
-        let lf = get_lf(item)?;
+        let lf = get_lf(&item)?;
         lfs.push(lf);
     }
 
@@ -270,7 +270,7 @@ pub fn datetime(
 
 #[pyfunction]
 pub fn concat_lf_diagonal(
-    lfs: &PyAny,
+    lfs: &Bound<'_, PyAny>,
     rechunk: bool,
     parallel: bool,
     to_supertypes: bool,
@@ -280,7 +280,7 @@ pub fn concat_lf_diagonal(
     let lfs = iter
         .map(|item| {
             let item = item?;
-            get_lf(item)
+            get_lf(&item)
         })
         .collect::<PyResult<Vec<_>>>()?;
 
@@ -298,13 +298,13 @@ pub fn concat_lf_diagonal(
 }
 
 #[pyfunction]
-pub fn concat_lf_horizontal(lfs: &PyAny, parallel: bool) -> PyResult<PyLazyFrame> {
+pub fn concat_lf_horizontal(lfs: &Bound<'_, PyAny>, parallel: bool) -> PyResult<PyLazyFrame> {
     let iter = lfs.iter()?;
 
     let lfs = iter
         .map(|item| {
             let item = item?;
-            get_lf(item)
+            get_lf(&item)
         })
         .collect::<PyResult<Vec<_>>>()?;
 
@@ -392,7 +392,7 @@ pub fn nth(n: i64) -> PyExpr {
 }
 
 #[pyfunction]
-pub fn lit(value: &PyAny, allow_object: bool) -> PyResult<PyExpr> {
+pub fn lit(value: &Bound<'_, PyAny>, allow_object: bool) -> PyResult<PyExpr> {
     if value.is_instance_of::<PyBool>() {
         let val = value.extract::<bool>().unwrap();
         Ok(dsl::lit(val).into())
@@ -406,12 +406,7 @@ pub fn lit(value: &PyAny, allow_object: bool) -> PyResult<PyExpr> {
         let val = float.extract::<f64>().unwrap();
         Ok(Expr::Literal(LiteralValue::Float(val)).into())
     } else if let Ok(pystr) = value.downcast::<PyString>() {
-        Ok(dsl::lit(
-            pystr
-                .to_str()
-                .expect("could not transform Python string to Rust Unicode"),
-        )
-        .into())
+        Ok(dsl::lit(pystr.to_string()).into())
     } else if let Ok(series) = value.extract::<PySeries>() {
         Ok(dsl::lit(series.series).into())
     } else if value.is_none() {

diff --git a/py-polars/src/functions/range.rs b/py-polars/src/functions/range.rs
@@ -17,9 +17,9 @@ pub fn int_range(start: PyExpr, end: PyExpr, step: i64, dtype: Wrap<DataType>) -
 /// Eager version of `int_range` to avoid overhead from the expression engine.
 #[pyfunction]
 pub fn eager_int_range(
-    lower: &PyAny,
-    upper: &PyAny,
-    step: &PyAny,
+    lower: &Bound<'_, PyAny>,
+    upper: &Bound<'_, PyAny>,
+    step: &Bound<'_, PyAny>,
     dtype: Wrap<DataType>,
 ) -> PyResult<PySeries> {
     let dtype = dtype.0;