refactor PyDict::get_item[_with_error] implementations

PyO3 · Jun 25, 2023 · 20f909c · 20f909c
1 parent db91642
commit 20f909c
Showing 1 changed file with 23 additions and 15 deletions.
diff --git a/src/types/dict.rs b/src/types/dict.rs
@@ -2,10 +2,9 @@ use super::PyMapping;
 use crate::err::{self, PyErr, PyResult};
 use crate::ffi::Py_ssize_t;
 use crate::types::{PyAny, PyList};
-use crate::{ffi, AsPyPointer, Python, ToPyObject};
 #[cfg(not(PyPy))]
-use crate::{IntoPyPointer, PyObject};
-use std::ptr::NonNull;
+use crate::IntoPyPointer;
+use crate::{ffi, AsPyPointer, PyObject, Python, ToPyObject};
 
 /// Represents a Python `dict`.
 #[repr(transparent)]
@@ -143,12 +142,16 @@ impl PyDict {
     where
         K: ToPyObject,
     {
+        self.get_item_impl(key.to_object(self.py()))
+    }
+
+    fn get_item_impl(&self, key: PyObject) -> Option<&PyAny> {
+        let py = self.py();
         unsafe {
-            let ptr = ffi::PyDict_GetItem(self.as_ptr(), key.to_object(self.py()).as_ptr());
-            NonNull::new(ptr).map(|p| {
-                // PyDict_GetItem return s borrowed ptr, must make it owned for safety (see #890).
-                self.py().from_owned_ptr(ffi::_Py_NewRef(p.as_ptr()))
-            })
+            let ptr = ffi::PyDict_GetItem(self.as_ptr(), key.as_ptr());
+            // PyDict_GetItem returns a borrowed ptr, must make it owned for safety (see #890).
+            // PyObject::from_borrowed_ptr_or_opt will take ownership in this way.
+            PyObject::from_borrowed_ptr_or_opt(py, ptr).map(|pyobject| pyobject.into_ref(py))
         }
     }
 
@@ -161,14 +164,19 @@ impl PyDict {
     where
         K: ToPyObject,
     {
-        unsafe {
-            let ptr =
-                ffi::PyDict_GetItemWithError(self.as_ptr(), key.to_object(self.py()).as_ptr());
-            if !ffi::PyErr_Occurred().is_null() {
-                return Err(PyErr::fetch(self.py()));
-            }
+        self.get_item_with_error_impl(key.to_object(self.py()))
+    }
 
-            Ok(NonNull::new(ptr).map(|p| self.py().from_owned_ptr(ffi::_Py_NewRef(p.as_ptr()))))
+    fn get_item_with_error_impl(&self, key: PyObject) -> PyResult<Option<&PyAny>> {
+        let py = self.py();
+        unsafe {
+            let ptr = ffi::PyDict_GetItemWithError(self.as_ptr(), key.as_ptr());
+            // PyDict_GetItemWithError returns a borrowed ptr, must make it owned for safety (see #890).
+            // PyObject::from_borrowed_ptr_or_opt will take ownership in this way.
+            PyObject::from_borrowed_ptr_or_opt(py, ptr)
+                .map(|pyobject| Ok(pyobject.into_ref(py)))
+                .or_else(|| PyErr::take(py).map(Err))
+                .transpose()
         }
     }