Merge pull request #1825 from PyO3/index_trait

2021-08-24 09:34:42 +02:00 · 2021-08-24 09:34:42 +02:00 · 50cd4c60ca
parent cfd6a5bc35 a0f7488e1a
commit 50cd4c60ca
5 changed files with 164 additions and 40 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.

 - Add `PyList::get_item_unchecked()` and `PyTuple::get_item_unchecked()` to get items without bounds checks. [#1733](https://github.com/PyO3/pyo3/pull/1733)
 - Add `PyAny::py()` as a convenience for `PyNativeType::py()`. [#1751](https://github.com/PyO3/pyo3/pull/1751)
+- Add implementation of `std::ops::Index` for `PyList`, `PyTuple` and `PySequence`. [#1825](https://github.com/PyO3/pyo3/pull/1825)

 ### Changed

--- a/guide/src/migration.md
+++ b/guide/src/migration.md
@ -3,6 +3,25 @@
 This guide can help you upgrade code through breaking changes from one PyO3 version to the next.
 For a detailed list of all changes, see the [CHANGELOG](changelog.md).

+## from 0.14.* to 0.15
+
+### Changes in sequence indexing
+
+For all types that take sequence indices (`PyList`, `PyTuple` and `PySequence`),
+the API has been made consistent to only take `usize` indices, for consistency
+with Rust's indexing conventions.  Negative indices, which were only
+sporadically supported even in APIs that took `isize`, now aren't supported
+anywhere.
+
+Further, the `get_item` methods now always return a `PyResult` instead of
+panicking on invalid indices.  The `Index` trait has been implemented instead,
+and provides the same panic behavior as on Rust vectors.
+
+Note that *slice* indices (accepted by `PySequence::get_slice` and other) still
+inherit the Python behavior of clamping the indices to the actual length, and
+not panicking/returning an error on out of range indices.
+
+
 ## from 0.13.* to 0.14

 ### `auto-initialize` feature is now opt-in
--- a/src/types/list.rs
+++ b/src/types/list.rs
@ -8,6 +8,7 @@ use crate::internal_tricks::get_ssize_index;
 use crate::{
    AsPyPointer, IntoPy, IntoPyPointer, PyAny, PyObject, Python, ToBorrowedObject, ToPyObject,
 };
+use std::ops::Index;

 /// Represents a Python `list`.
 #[repr(transparent)]
@ -176,6 +177,20 @@ impl PyList {
    }
 }

+impl Index<usize> for PyList {
+    type Output = PyAny;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        self.get_item(index).unwrap_or_else(|_| {
+            panic!(
+                "index {} out of range for list of length {}",
+                index,
+                self.len()
+            );
+        })
+    }
+}
+
 /// Used by `PyList::iter()`.
 pub struct PyListIterator<'a> {
    list: &'a PyList,
@ -189,7 +204,7 @@ impl<'a> Iterator for PyListIterator<'a> {
    fn next(&mut self) -> Option<&'a PyAny> {
        if self.index < self.list.len() {
            #[cfg(any(Py_LIMITED_API, PyPy))]
-            let item = self.list.get_item(self.index).expect("tuple.get failed");
+            let item = self.list.get_item(self.index).expect("list.get failed");
            #[cfg(not(any(Py_LIMITED_API, PyPy)))]
            let item = unsafe { self.list.get_item_unchecked(self.index) };
            self.index += 1;
@ -256,10 +271,10 @@ mod tests {
    fn test_new() {
        Python::with_gil(|py| {
            let list = PyList::new(py, &[2, 3, 5, 7]);
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(5, list.get_item(2).unwrap().extract::<i32>().unwrap());
-            assert_eq!(7, list.get_item(3).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
+            assert_eq!(5, list[2].extract::<i32>().unwrap());
+            assert_eq!(7, list[3].extract::<i32>().unwrap());
        });
    }

@ -299,9 +314,9 @@ mod tests {
            let list = PyList::new(py, &[2, 3, 5, 7]);
            let val = 42i32.to_object(py);
            let val2 = 42i32.to_object(py);
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
            list.set_item(0, val).unwrap();
-            assert_eq!(42, list.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(42, list[0].extract::<i32>().unwrap());
            assert!(list.set_item(10, val2).is_err());
        });
    }
@ -331,13 +346,13 @@ mod tests {
            let val = 42i32.to_object(py);
            let val2 = 43i32.to_object(py);
            assert_eq!(4, list.len());
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
            list.insert(0, val).unwrap();
            list.insert(1000, val2).unwrap();
            assert_eq!(6, list.len());
-            assert_eq!(42, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(2, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(43, list.get_item(5).unwrap().extract::<i32>().unwrap());
+            assert_eq!(42, list[0].extract::<i32>().unwrap());
+            assert_eq!(2, list[1].extract::<i32>().unwrap());
+            assert_eq!(43, list[5].extract::<i32>().unwrap());
        });
    }

@ -362,8 +377,8 @@ mod tests {
        Python::with_gil(|py| {
            let list = PyList::new(py, &[2]);
            list.append(3).unwrap();
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(1).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
        });
    }

@ -440,15 +455,15 @@ mod tests {
        Python::with_gil(|py| {
            let v = vec![7, 3, 2, 5];
            let list = PyList::new(py, &v);
-            assert_eq!(7, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(2, list.get_item(2).unwrap().extract::<i32>().unwrap());
-            assert_eq!(5, list.get_item(3).unwrap().extract::<i32>().unwrap());
+            assert_eq!(7, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
+            assert_eq!(2, list[2].extract::<i32>().unwrap());
+            assert_eq!(5, list[3].extract::<i32>().unwrap());
            list.sort().unwrap();
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(5, list.get_item(2).unwrap().extract::<i32>().unwrap());
-            assert_eq!(7, list.get_item(3).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
+            assert_eq!(5, list[2].extract::<i32>().unwrap());
+            assert_eq!(7, list[3].extract::<i32>().unwrap());
        });
    }

@ -457,15 +472,15 @@ mod tests {
        Python::with_gil(|py| {
            let v = vec![2, 3, 5, 7];
            let list = PyList::new(py, &v);
-            assert_eq!(2, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(5, list.get_item(2).unwrap().extract::<i32>().unwrap());
-            assert_eq!(7, list.get_item(3).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
+            assert_eq!(5, list[2].extract::<i32>().unwrap());
+            assert_eq!(7, list[3].extract::<i32>().unwrap());
            list.reverse().unwrap();
-            assert_eq!(7, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(5, list.get_item(1).unwrap().extract::<i32>().unwrap());
-            assert_eq!(3, list.get_item(2).unwrap().extract::<i32>().unwrap());
-            assert_eq!(2, list.get_item(3).unwrap().extract::<i32>().unwrap());
+            assert_eq!(7, list[0].extract::<i32>().unwrap());
+            assert_eq!(5, list[1].extract::<i32>().unwrap());
+            assert_eq!(3, list[2].extract::<i32>().unwrap());
+            assert_eq!(2, list[3].extract::<i32>().unwrap());
        });
    }

@ -474,8 +489,8 @@ mod tests {
        Python::with_gil(|py| {
            let array: PyObject = [1, 2].into_py(py);
            let list = <PyList as PyTryFrom>::try_from(array.as_ref(py)).unwrap();
-            assert_eq!(1, list.get_item(0).unwrap().extract::<i32>().unwrap());
-            assert_eq!(2, list.get_item(1).unwrap().extract::<i32>().unwrap());
+            assert_eq!(1, list[0].extract::<i32>().unwrap());
+            assert_eq!(2, list[1].extract::<i32>().unwrap());
        });
    }

@ -510,4 +525,23 @@ mod tests {
            assert_eq!(obj.extract::<i32>().unwrap(), 2);
        });
    }
+
+    #[test]
+    fn test_list_index_trait() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            assert_eq!(2, list[0].extract::<i32>().unwrap());
+            assert_eq!(3, list[1].extract::<i32>().unwrap());
+            assert_eq!(5, list[2].extract::<i32>().unwrap());
+        });
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_list_index_trait_panic() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            let _ = &list[7];
+        });
+    }
 }
--- a/src/types/sequence.rs
+++ b/src/types/sequence.rs
@ -6,6 +6,7 @@ use crate::internal_tricks::get_ssize_index;
 use crate::types::{PyAny, PyList, PyTuple};
 use crate::AsPyPointer;
 use crate::{FromPyObject, PyTryFrom, ToBorrowedObject};
+use std::ops::Index;

 /// Represents a reference to a Python object supporting the sequence protocol.
 #[repr(transparent)]
@ -268,6 +269,16 @@ impl PySequence {
    }
 }

+impl Index<usize> for PySequence {
+    type Output = PyAny;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        self.get_item(index).unwrap_or_else(|_| {
+            panic!("index {} out of range for sequence", index);
+        })
+    }
+}
+
 impl<'a, T> FromPyObject<'a> for Vec<T>
 where
    T: FromPyObject<'a>,
@ -430,6 +441,29 @@ mod tests {
        });
    }

+    #[test]
+    fn test_seq_index_trait() {
+        Python::with_gil(|py| {
+            let v: Vec<i32> = vec![1, 1, 2];
+            let ob = v.to_object(py);
+            let seq = ob.cast_as::<PySequence>(py).unwrap();
+            assert_eq!(1, seq[0].extract::<i32>().unwrap());
+            assert_eq!(1, seq[1].extract::<i32>().unwrap());
+            assert_eq!(2, seq[2].extract::<i32>().unwrap());
+        });
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_seq_index_trait_panic() {
+        Python::with_gil(|py| {
+            let v: Vec<i32> = vec![1, 1, 2];
+            let ob = v.to_object(py);
+            let seq = ob.cast_as::<PySequence>(py).unwrap();
+            let _ = &seq[7];
+        });
+    }
+
    #[test]
    fn test_seq_del_item() {
        Python::with_gil(|py| {
@ -437,17 +471,17 @@ mod tests {
            let ob = v.to_object(py);
            let seq = ob.cast_as::<PySequence>(py).unwrap();
            assert!(seq.del_item(10).is_err());
-            assert_eq!(1, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(1, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
-            assert_eq!(1, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(1, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
-            assert_eq!(2, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
-            assert_eq!(3, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(3, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
-            assert_eq!(5, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(5, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
-            assert_eq!(8, seq.get_item(0).unwrap().extract::<i32>().unwrap());
+            assert_eq!(8, seq[0].extract::<i32>().unwrap());
            assert!(seq.del_item(0).is_ok());
            assert_eq!(0, seq.len().unwrap());
            assert!(seq.del_item(0).is_err());
@ -460,9 +494,9 @@ mod tests {
            let v: Vec<i32> = vec![1, 2];
            let ob = v.to_object(py);
            let seq = ob.cast_as::<PySequence>(py).unwrap();
-            assert_eq!(2, seq.get_item(1).unwrap().extract::<i32>().unwrap());
+            assert_eq!(2, seq[1].extract::<i32>().unwrap());
            assert!(seq.set_item(1, 10).is_ok());
-            assert_eq!(10, seq.get_item(1).unwrap().extract::<i32>().unwrap());
+            assert_eq!(10, seq[1].extract::<i32>().unwrap());
        });
    }

@ -475,7 +509,7 @@ mod tests {
            let ob = v.to_object(py);
            let seq = ob.cast_as::<PySequence>(py).unwrap();
            assert!(seq.set_item(1, &obj).is_ok());
-            assert!(seq.get_item(1).unwrap().as_ptr() == obj.as_ptr());
+            assert!(seq[1].as_ptr() == obj.as_ptr());
        });

        Python::with_gil(|py| {
--- a/src/types/tuple.rs
+++ b/src/types/tuple.rs
@ -6,6 +6,7 @@ use crate::{
    exceptions, AsPyPointer, FromPyObject, IntoPy, IntoPyPointer, Py, PyAny, PyErr, PyObject,
    PyResult, PyTryFrom, Python, ToPyObject,
 };
+use std::ops::Index;

 /// Represents a Python `tuple` object.
 ///
@ -139,6 +140,20 @@ impl PyTuple {
    }
 }

+impl Index<usize> for PyTuple {
+    type Output = PyAny;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        self.get_item(index).unwrap_or_else(|_| {
+            panic!(
+                "index {} out of range for tuple of length {}",
+                index,
+                self.len()
+            );
+        })
+    }
+}
+
 /// Used by `PyTuple::iter()`.
 pub struct PyTupleIterator<'a> {
    tuple: &'a PyTuple,
@ -537,4 +552,25 @@ mod tests {
            assert_eq!(obj.extract::<i32>().unwrap(), 1);
        });
    }
+
+    #[test]
+    fn test_tuple_index_trait() {
+        Python::with_gil(|py| {
+            let ob = (1, 2, 3).to_object(py);
+            let tuple = <PyTuple as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
+            assert_eq!(1, tuple[0].extract::<i32>().unwrap());
+            assert_eq!(2, tuple[1].extract::<i32>().unwrap());
+            assert_eq!(3, tuple[2].extract::<i32>().unwrap());
+        });
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_tuple_index_trait_panic() {
+        Python::with_gil(|py| {
+            let ob = (1, 2, 3).to_object(py);
+            let tuple = <PyTuple as PyTryFrom>::try_from(ob.as_ref(py)).unwrap();
+            let _ = &tuple[7];
+        });
+    }
 }