Merge pull request #1682 from moriyoshi/allow-as-ref-pysequence

Allow calling Py::as_ref() and Py::into_ref() against PySequence, PyIterator and PyMapping.

CHANGELOG.md

@@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add implementations for `Py::as_ref()` and `Py::into_ref()` for `Py<PySequence>`, `Py<PyIterator>` and `Py<PyMapping>`. [#1682](https://github.com/PyO3/pyo3/pull/1682)
+
 ### Changed
 
 - `#[classattr]` constants with a known magic method name (which is lowercase) no longer trigger lint warnings expecting constants to be uppercase. [#1969](https://github.com/PyO3/pyo3/pull/1969)

src/types/iterator.rs

@@ -2,7 +2,7 @@
 //
 // based on Daniel Grunwald's https://github.com/dgrunwald/rust-cpython
 
-use crate::{ffi, AsPyPointer, PyAny, PyErr, PyResult, Python};
+use crate::{ffi, AsPyPointer, IntoPyPointer, Py, PyAny, PyErr, PyNativeType, PyResult, Python};
 #[cfg(any(not(Py_LIMITED_API), Py_3_8))]
 use crate::{PyDowncastError, PyTryFrom};
 
@@ -88,15 +88,31 @@ impl<'v> PyTryFrom<'v> for PyIterator {
     }
 }
 
+impl Py<PyIterator> {
+    /// Borrows a GIL-bound reference to the PyIterator. By binding to the GIL lifetime, this
+    /// allows the GIL-bound reference to not require `Python` for any of its methods.
+    pub fn as_ref<'py>(&'py self, _py: Python<'py>) -> &'py PyIterator {
+        let any = self.as_ptr() as *const PyAny;
+        unsafe { PyNativeType::unchecked_downcast(&*any) }
+    }
+
+    /// Similar to [`as_ref`](#method.as_ref), and also consumes this `Py` and registers the
+    /// Python object reference in PyO3's object storage. The reference count for the Python
+    /// object will not be decreased until the GIL lifetime ends.
+    pub fn into_ref(self, py: Python) -> &PyIterator {
+        unsafe { py.from_owned_ptr(self.into_ptr()) }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::PyIterator;
     use crate::exceptions::PyTypeError;
     use crate::gil::GILPool;
     use crate::types::{PyDict, PyList};
     #[cfg(any(not(Py_LIMITED_API), Py_3_8))]
-    use crate::{Py, PyAny, PyTryFrom};
-    use crate::{Python, ToPyObject};
+    use crate::PyTryFrom;
+    use crate::{Py, PyAny, Python, ToPyObject};
     use indoc::indoc;
 
     #[test]
@@ -203,4 +219,26 @@ mod tests {
             assert_eq!(obj, iter.into());
         });
     }
+
+    #[test]
+    fn test_as_ref() {
+        Python::with_gil(|py| {
+            let iter: Py<PyIterator> = PyAny::iter(PyList::empty(py)).unwrap().into();
+            let mut iter_ref: &PyIterator = iter.as_ref(py);
+            assert!(iter_ref.next().is_none());
+        })
+    }
+
+    #[test]
+    fn test_into_ref() {
+        Python::with_gil(|py| {
+            let bare_iter = PyAny::iter(PyList::empty(py)).unwrap();
+            assert_eq!(bare_iter.get_refcnt(), 1);
+            let iter: Py<PyIterator> = bare_iter.into();
+            assert_eq!(bare_iter.get_refcnt(), 2);
+            let mut iter_ref = iter.into_ref(py);
+            assert!(iter_ref.next().is_none());
+            assert_eq!(iter_ref.get_refcnt(), 2);
+        })
+    }
 }

src/types/mapping.rs

@@ -2,9 +2,10 @@
 
 use crate::err::{PyDowncastError, PyErr, PyResult};
 use crate::types::{PyAny, PySequence};
-use crate::AsPyPointer;
-use crate::{ffi, ToPyObject};
-use crate::{PyTryFrom, ToBorrowedObject};
+use crate::{
+    ffi, AsPyPointer, IntoPyPointer, Py, PyNativeType, PyTryFrom, Python, ToBorrowedObject,
+    ToPyObject,
+};
 
 /// Represents a reference to a Python object supporting the mapping protocol.
 #[repr(transparent)]
@@ -120,11 +121,31 @@ impl<'v> PyTryFrom<'v> for PyMapping {
     }
 }
 
+impl Py<PyMapping> {
+    /// Borrows a GIL-bound reference to the PyMapping. By binding to the GIL lifetime, this
+    /// allows the GIL-bound reference to not require `Python` for any of its methods.
+    pub fn as_ref<'py>(&'py self, _py: Python<'py>) -> &'py PyMapping {
+        let any = self.as_ptr() as *const PyAny;
+        unsafe { PyNativeType::unchecked_downcast(&*any) }
+    }
+
+    /// Similar to [`as_ref`](#method.as_ref), and also consumes this `Py` and registers the
+    /// Python object reference in PyO3's object storage. The reference count for the Python
+    /// object will not be decreased until the GIL lifetime ends.
+    pub fn into_ref(self, py: Python) -> &PyMapping {
+        unsafe { py.from_owned_ptr(self.into_ptr()) }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
 
-    use crate::{exceptions::PyKeyError, types::PyTuple, Python};
+    use crate::{
+        exceptions::PyKeyError,
+        types::{PyDict, PyTuple},
+        Python,
+    };
 
     use super::*;
 
@@ -256,4 +277,26 @@ mod tests {
             assert_eq!(32 + 42 + 123, values_sum);
         });
     }
+
+    #[test]
+    fn test_as_ref() {
+        Python::with_gil(|py| {
+            let mapping: Py<PyMapping> = PyDict::new(py).as_mapping().into();
+            let mapping_ref: &PyMapping = mapping.as_ref(py);
+            assert_eq!(mapping_ref.len().unwrap(), 0);
+        })
+    }
+
+    #[test]
+    fn test_into_ref() {
+        Python::with_gil(|py| {
+            let bare_mapping = PyDict::new(py).as_mapping();
+            assert_eq!(bare_mapping.get_refcnt(), 1);
+            let mapping: Py<PyMapping> = bare_mapping.into();
+            assert_eq!(bare_mapping.get_refcnt(), 2);
+            let mapping_ref = mapping.into_ref(py);
+            assert_eq!(mapping_ref.len().unwrap(), 0);
+            assert_eq!(mapping_ref.get_refcnt(), 2);
+        })
+    }
 }

src/types/sequence.rs

@@ -1,10 +1,10 @@
 // Copyright (c) 2017-present PyO3 Project and Contributors
 
 use crate::err::{self, PyDowncastError, PyErr, PyResult};
-use crate::ffi;
 use crate::internal_tricks::get_ssize_index;
 use crate::types::{PyAny, PyList, PyTuple};
-use crate::AsPyPointer;
+use crate::{ffi, PyNativeType};
+use crate::{AsPyPointer, IntoPyPointer, Py, Python};
 use crate::{FromPyObject, PyTryFrom, ToBorrowedObject};
 
 /// Represents a reference to a Python object supporting the sequence protocol.
@@ -335,12 +335,36 @@ impl<'v> PyTryFrom<'v> for PySequence {
     }
 }
 
+impl Py<PySequence> {
+    /// Borrows a GIL-bound reference to the PySequence. By binding to the GIL lifetime, this
+    /// allows the GIL-bound reference to not require `Python` for any of its methods.
+    ///
+    /// ```
+    /// # use pyo3::prelude::*;
+    /// # use pyo3::types::{PyList, PySequence};
+    /// # Python::with_gil(|py| {
+    /// let seq: Py<PySequence> = PyList::empty(py).as_sequence().into();
+    /// let seq: &PySequence = seq.as_ref(py);
+    /// assert_eq!(seq.len().unwrap(), 0);
+    /// # });
+    /// ```
+    pub fn as_ref<'py>(&'py self, _py: Python<'py>) -> &'py PySequence {
+        let any = self.as_ptr() as *const PyAny;
+        unsafe { PyNativeType::unchecked_downcast(&*any) }
+    }
+
+    /// Similar to [`as_ref`](#method.as_ref), and also consumes this `Py` and registers the
+    /// Python object reference in PyO3's object storage. The reference count for the Python
+    /// object will not be decreased until the GIL lifetime ends.
+    pub fn into_ref(self, py: Python) -> &PySequence {
+        unsafe { py.from_owned_ptr(self.into_ptr()) }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::types::{PyList, PySequence};
-    use crate::AsPyPointer;
-    use crate::Python;
-    use crate::{PyObject, PyTryFrom, ToPyObject};
+    use crate::{AsPyPointer, Py, PyObject, PyTryFrom, Python, ToPyObject};
 
     fn get_object() -> PyObject {
         // Convenience function for getting a single unique object
@@ -818,4 +842,26 @@ mod tests {
             assert!(seq_from.list().is_ok());
         });
     }
+
+    #[test]
+    fn test_as_ref() {
+        Python::with_gil(|py| {
+            let seq: Py<PySequence> = PyList::empty(py).as_sequence().into();
+            let seq_ref: &PySequence = seq.as_ref(py);
+            assert_eq!(seq_ref.len().unwrap(), 0);
+        })
+    }
+
+    #[test]
+    fn test_into_ref() {
+        Python::with_gil(|py| {
+            let bare_seq = PyList::empty(py).as_sequence();
+            assert_eq!(bare_seq.get_refcnt(), 1);
+            let seq: Py<PySequence> = bare_seq.into();
+            assert_eq!(bare_seq.get_refcnt(), 2);
+            let seq_ref = seq.into_ref(py);
+            assert_eq!(seq_ref.len().unwrap(), 0);
+            assert_eq!(seq_ref.get_refcnt(), 2);
+        })
+    }
 }