Merge pull request #1027 from noam93k/bugfix/convert-bigint-index

Use the result of __index__ when converting to Rust BigInts.
PyO3 · Jul 9, 2020 · bd97b25 · bd97b25
2 parents c00080e + e5959f2
commit bd97b25
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Changed
 - Change FFI definitions `Py_SetProgramName` and `Py_SetPythonHome` to take `*const` argument instead of `*mut`. [#1021](https://github.com/PyO3/pyo3/pull/1021)
 
+### Fixed
+- Conversion from types with an `__index__` method to Rust BigInts. [#1027](https://github.com/PyO3/pyo3/pull/1027)
+
 ## [0.11.1] - 2020-06-30
 ### Added
 - `#[pyclass(unsendable)]`. [#1009](https://github.com/PyO3/pyo3/pull/1009)

diff --git a/src/types/num.rs b/src/types/num.rs
@@ -279,24 +279,28 @@ mod bigint_conversion {
             }
             impl<'source> FromPyObject<'source> for $rust_ty {
                 fn extract(ob: &'source PyAny) -> PyResult<$rust_ty> {
+                    use crate::instance::AsPyRef;
+                    let py = ob.py();
                     unsafe {
                         let num = ffi::PyNumber_Index(ob.as_ptr());
                         if num.is_null() {
-                            return Err(PyErr::fetch(ob.py()));
+                            return Err(PyErr::fetch(py));
                         }
                         let n_bits = ffi::_PyLong_NumBits(num);
                         let n_bytes = if n_bits < 0 {
-                            return Err(PyErr::fetch(ob.py()));
+                            return Err(PyErr::fetch(py));
                         } else if n_bits == 0 {
                             0
                         } else {
                             (n_bits as usize - 1 + $is_signed) / 8 + 1
                         };
+                        let ob = PyObject::from_owned_ptr(py, num);
                         if n_bytes <= 128 {
-                            extract_small(ob, n_bytes, $is_signed)
+                            extract_small(ob.as_ref(py), n_bytes, $is_signed)
                                 .map(|b| $from_bytes(&b[..n_bytes]))
                         } else {
-                            extract_large(ob, n_bytes, $is_signed).map(|b| $from_bytes(&b))
+                            extract_large(ob.as_ref(py), n_bytes, $is_signed)
+                                .map(|b| $from_bytes(&b))
                         }
                     }
                 }
@@ -393,6 +397,30 @@ mod bigint_conversion {
             assert_eq!(rs_result, py_result);
         }
 
+        fn python_index_class(py: Python) -> &PyModule {
+            let index_code = indoc!(
+                r#"
+                class C:
+                    def __init__(self, x):
+                        self.x = x
+                    def __index__(self):
+                        return self.x
+                "#
+            );
+            PyModule::from_code(py, index_code, "index.py", "index").unwrap()
+        }
+
+        #[test]
+        fn convert_index_class() {
+            let gil = Python::acquire_gil();
+            let py = gil.python();
+            let index = python_index_class(py);
+            let locals = PyDict::new(py);
+            locals.set_item("index", index).unwrap();
+            let ob = py.eval("index.C(10)", None, Some(locals)).unwrap();
+            let _: BigInt = FromPyObject::extract(ob).unwrap();
+        }
+
         #[test]
         fn handle_zero() {
             let gil = Python::acquire_gil();