list,tuple,sequence: add slice indexing

CHANGELOG.md

@@ -12,19 +12,20 @@ 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)
+- Add implementation of `std::ops::Index<usize>` for `PyList`, `PyTuple` and `PySequence`. [#1825](https://github.com/PyO3/pyo3/pull/1825)
+- Add range indexing implementations of `std::ops::Index` for `PyList`, `PyTuple` and `PySequence`. [#1829](https://github.com/PyO3/pyo3/pull/1829)
 
 ### Changed
 
 - Change `PyErr::fetch()` to return `Option<PyErr>`. [#1717](https://github.com/PyO3/pyo3/pull/1717)
 - `PyList`, `PyTuple` and `PySequence`'s APIs now accepts only `usize` indices instead of `isize`.
   [#1733](https://github.com/PyO3/pyo3/pull/1733), [#1802](https://github.com/PyO3/pyo3/pull/1802),
   [#1803](https://github.com/PyO3/pyo3/pull/1803)
-- `PyList` and `PyTuple`'s `get_item` now return `PyResult<&PyAny>` instead of panicking. [#1733](https://github.com/PyO3/pyo3/pull/1733)
-- `PySequence`'s `in_place_repeat` and `in_place_concat` now return the result instead of `()`, which is needed
-  in case of immutable sequences. [#1803](https://github.com/PyO3/pyo3/pull/1803)
+- `PyList::get_item` and `PyTuple::get_item` now return `PyResult<&PyAny>` instead of panicking. [#1733](https://github.com/PyO3/pyo3/pull/1733)
+- `PySequence::in_place_repeat` and `PySequence::in_place_concat` now return `PyResult<&PySequence>` instead of `PyResult<()>`, which is needed in case of immutable sequences such as tuples. [#1803](https://github.com/PyO3/pyo3/pull/1803)
+- `PySequence::get_slice` now returns `PyResult<&PySequence>` instead of `PyResult<&PyAny>`. [#1829](https://github.com/PyO3/pyo3/pull/1829)
 - Deprecate `PyTuple::split_from`. [#1804](https://github.com/PyO3/pyo3/pull/1804)
-- Deprecate `PyTuple::slice`, new method `get_slice` added with proper index types. [#1828](https://github.com/PyO3/pyo3/pull/1828)
+- Deprecate `PyTuple::slice`, new method `PyTuple::get_slice` added with `usize` indices. [#1828](https://github.com/PyO3/pyo3/pull/1828)
 
 ## [0.14.3] - 2021-08-22
 

build.rs

@@ -137,6 +137,11 @@ fn configure_pyo3() -> Result<()> {
 
     let rustc_minor_version = rustc_minor_version().unwrap_or(0);
 
+    // Enable use of #[track_caller] on Rust 1.46 and greater
+    if rustc_minor_version >= 46 {
+        println!("cargo:rustc-cfg=track_caller");
+    }
+
     // Enable use of const generics on Rust 1.51 and greater
     if rustc_minor_version >= 51 {
         println!("cargo:rustc-cfg=min_const_generics");

guide/src/migration.md

@@ -21,6 +21,18 @@ 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.
 
+An additional advantage of using Rust's indexing conventions for these types is
+that these types can now also support Rust's indexing operators as part of a
+consistent API:
+
+```rust
+use pyo3::{Python, types::PyList};
+
+Python::with_gil(|py| {
+    let list = PyList::new(py, &[1, 2, 3]);
+    assert_eq!(list[0..2].to_string(), "[1, 2]");
+});
+```
 
 ## from 0.13.* to 0.14
 

src/ffi/abstract_.rs

@@ -251,9 +251,9 @@ extern "C" {
     pub fn PySequence_List(o: *mut PyObject) -> *mut PyObject;
     #[cfg_attr(PyPy, link_name = "PyPySequence_Fast")]
     pub fn PySequence_Fast(o: *mut PyObject, m: *const c_char) -> *mut PyObject;
-    // skipped PySequenc_Fast_GET_SIZE
-    // skipped PySequenc_Fast_GET_ITEM
-    // skipped PySequenc_Fast_GET_ITEMS
+    // skipped PySequence_Fast_GET_SIZE
+    // skipped PySequence_Fast_GET_ITEM
+    // skipped PySequence_Fast_GET_ITEMS
     pub fn PySequence_Count(o: *mut PyObject, value: *mut PyObject) -> Py_ssize_t;
     #[cfg_attr(PyPy, link_name = "PyPySequence_Contains")]
     pub fn PySequence_Contains(seq: *mut PyObject, ob: *mut PyObject) -> c_int;

src/internal_tricks.rs

@@ -58,3 +58,146 @@ pub(crate) fn extract_cstr_or_leak_cstring(
 pub(crate) fn get_ssize_index(index: usize) -> Py_ssize_t {
     index.min(PY_SSIZE_T_MAX as usize) as Py_ssize_t
 }
+
+/// Implementations used for slice indexing PySequence, PyTuple, and PyList
+macro_rules! index_impls {
+    (
+        $ty:ty,
+        $ty_name:literal,
+        $len:expr,
+        $get_slice:expr $(,)?
+    ) => {
+        impl std::ops::Index<usize> for $ty {
+            // Always PyAny output (even if the slice operation returns something else)
+            type Output = PyAny;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(&self, index: usize) -> &Self::Output {
+                self.get_item(index).unwrap_or_else(|_| {
+                    crate::internal_tricks::index_len_fail(index, $ty_name, $len(self))
+                })
+            }
+        }
+
+        impl std::ops::Index<std::ops::Range<usize>> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(
+                &self,
+                std::ops::Range { start, end }: std::ops::Range<usize>,
+            ) -> &Self::Output {
+                let len = $len(self);
+                if start > len {
+                    crate::internal_tricks::slice_start_index_len_fail(start, $ty_name, len)
+                } else if end > len {
+                    crate::internal_tricks::slice_end_index_len_fail(end, $ty_name, len)
+                } else if start > end {
+                    crate::internal_tricks::slice_index_order_fail(start, end)
+                } else {
+                    $get_slice(self, start, end)
+                }
+            }
+        }
+
+        impl std::ops::Index<std::ops::RangeFrom<usize>> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(
+                &self,
+                std::ops::RangeFrom { start }: std::ops::RangeFrom<usize>,
+            ) -> &Self::Output {
+                let len = $len(self);
+                if start > len {
+                    crate::internal_tricks::slice_start_index_len_fail(start, $ty_name, len)
+                } else {
+                    $get_slice(self, start, len)
+                }
+            }
+        }
+
+        impl std::ops::Index<std::ops::RangeFull> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(&self, _: std::ops::RangeFull) -> &Self::Output {
+                let len = $len(self);
+                $get_slice(self, 0, len)
+            }
+        }
+
+        impl std::ops::Index<std::ops::RangeInclusive<usize>> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(&self, range: std::ops::RangeInclusive<usize>) -> &Self::Output {
+                let exclusive_end = range
+                    .end()
+                    .checked_add(1)
+                    .expect("range end exceeds Python limit");
+                &self[*range.start()..exclusive_end]
+            }
+        }
+
+        impl std::ops::Index<std::ops::RangeTo<usize>> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(&self, std::ops::RangeTo { end }: std::ops::RangeTo<usize>) -> &Self::Output {
+                &self[0..end]
+            }
+        }
+
+        impl std::ops::Index<std::ops::RangeToInclusive<usize>> for $ty {
+            type Output = $ty;
+
+            #[cfg_attr(track_caller, track_caller)]
+            fn index(
+                &self,
+                std::ops::RangeToInclusive { end }: std::ops::RangeToInclusive<usize>,
+            ) -> &Self::Output {
+                &self[0..=end]
+            }
+        }
+    };
+}
+
+// these error messages are shamelessly "borrowed" from std.
+
+#[inline(never)]
+#[cold]
+#[cfg_attr(track_caller, track_caller)]
+pub(crate) fn index_len_fail(index: usize, ty_name: &str, len: usize) -> ! {
+    panic!(
+        "index {} out of range for {} of length {}",
+        index, ty_name, len
+    );
+}
+
+#[inline(never)]
+#[cold]
+#[cfg_attr(track_caller, track_caller)]
+pub(crate) fn slice_start_index_len_fail(index: usize, ty_name: &str, len: usize) -> ! {
+    panic!(
+        "range start index {} out of range for {} of length {}",
+        index, ty_name, len
+    );
+}
+
+#[inline(never)]
+#[cold]
+#[cfg_attr(track_caller, track_caller)]
+pub(crate) fn slice_end_index_len_fail(index: usize, ty_name: &str, len: usize) -> ! {
+    panic!(
+        "range end index {} out of range for {} of length {}",
+        index, ty_name, len
+    );
+}
+
+#[inline(never)]
+#[cold]
+#[cfg_attr(track_caller, track_caller)]
+pub(crate) fn slice_index_order_fail(index: usize, end: usize) -> ! {
+    panic!("slice index starts at {} but ends at {}", index, end);
+}

src/types/list.rs

@@ -8,7 +8,6 @@ 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)]
@@ -177,19 +176,7 @@ 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()
-            );
-        })
-    }
-}
+index_impls!(PyList, "list", PyList::len, PyList::get_slice);
 
 /// Used by `PyList::iter()`.
 pub struct PyListIterator<'a> {
@@ -544,4 +531,56 @@ mod tests {
             let _ = &list[7];
         });
     }
+
+    #[test]
+    fn test_list_index_trait_ranges() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            assert_eq!(vec![3, 5], list[1..3].extract::<Vec<i32>>().unwrap());
+            assert_eq!(Vec::<i32>::new(), list[3..3].extract::<Vec<i32>>().unwrap());
+            assert_eq!(vec![3, 5], list[1..].extract::<Vec<i32>>().unwrap());
+            assert_eq!(Vec::<i32>::new(), list[3..].extract::<Vec<i32>>().unwrap());
+            assert_eq!(vec![2, 3, 5], list[..].extract::<Vec<i32>>().unwrap());
+            assert_eq!(vec![3, 5], list[1..=2].extract::<Vec<i32>>().unwrap());
+            assert_eq!(vec![2, 3], list[..2].extract::<Vec<i32>>().unwrap());
+            assert_eq!(vec![2, 3], list[..=1].extract::<Vec<i32>>().unwrap());
+        })
+    }
+
+    #[test]
+    #[should_panic = "range start index 5 out of range for list of length 3"]
+    fn test_list_index_trait_range_panic_start() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            list[5..10].extract::<Vec<i32>>().unwrap();
+        })
+    }
+
+    #[test]
+    #[should_panic = "range end index 10 out of range for list of length 3"]
+    fn test_list_index_trait_range_panic_end() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            list[1..10].extract::<Vec<i32>>().unwrap();
+        })
+    }
+
+    #[test]
+    #[should_panic = "slice index starts at 2 but ends at 1"]
+    fn test_list_index_trait_range_panic_wrong_order() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            #[allow(clippy::reversed_empty_ranges)]
+            list[2..1].extract::<Vec<i32>>().unwrap();
+        })
+    }
+
+    #[test]
+    #[should_panic = "range start index 8 out of range for list of length 3"]
+    fn test_list_index_trait_range_from_panic() {
+        Python::with_gil(|py| {
+            let list = PyList::new(py, &[2, 3, 5]);
+            list[8..].extract::<Vec<i32>>().unwrap();
+        })
+    }
 }