Skip to content

Commit

Permalink
list,tuple,sequence: add slice indexing
Browse files Browse the repository at this point in the history
  • Loading branch information
davidhewitt committed Aug 25, 2021
1 parent 2df0cb6 commit 21759b9
Show file tree
Hide file tree
Showing 6 changed files with 340 additions and 38 deletions.
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ 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

Expand Down
5 changes: 5 additions & 0 deletions build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
143 changes: 143 additions & 0 deletions src/internal_tricks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
$output:ty,
$ty_name:literal,
$len:expr,
$get_slice:expr $(,)?
) => {
impl 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));
unreachable!()
})
}
}

impl Index<Range<usize>> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, Range { start, end }: Range<usize>) -> &Self::Output {
let len = $len(self);
if start > len {
crate::internal_tricks::slice_start_index_len_fail(start, $ty_name, len);
unreachable!()
} else if end > len {
crate::internal_tricks::slice_end_index_len_fail(end, $ty_name, len);
unreachable!()
} else if start > end {
crate::internal_tricks::slice_index_order_fail(start, end);
unreachable!()
} else {
$get_slice(self, start, end)
}
}
}

impl Index<RangeFrom<usize>> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, RangeFrom { start }: RangeFrom<usize>) -> &Self::Output {
let len = $len(self);
if start > len {
crate::internal_tricks::slice_start_index_len_fail(start, $ty_name, len);
unreachable!()
} else {
$get_slice(self, start, len)
}
}
}

impl Index<RangeFull> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, _: RangeFull) -> &Self::Output {
let len = $len(self);
$get_slice(self, 0, len)
}
}

impl Index<RangeInclusive<usize>> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, range: RangeInclusive<usize>) -> &Self::Output {
let exclusive_end = range
.end()
.checked_add(1)
.expect("range end exceeds Python limit");
&self[*range.start()..exclusive_end]
}
}

impl Index<RangeTo<usize>> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, RangeTo { end }: RangeTo<usize>) -> &Self::Output {
&self[0..end]
}
}

impl Index<RangeToInclusive<usize>> for $ty {
type Output = $output;

#[cfg_attr(track_caller, track_caller)]
fn index(&self, RangeToInclusive { end }: RangeToInclusive<usize>) -> &Self::Output {
&self[0..=end]
}
}
};
}

// nb these would all return ! (the never type) if it was stable, which would remove need for the
// unreachable!() uses above

// 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);
}
68 changes: 54 additions & 14 deletions src/types/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::internal_tricks::get_ssize_index;
use crate::{
AsPyPointer, IntoPy, IntoPyPointer, PyAny, PyObject, Python, ToBorrowedObject, ToPyObject,
};
use std::ops::Index;
use std::ops::{Index, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};

/// Represents a Python `list`.
#[repr(transparent)]
Expand Down Expand Up @@ -177,19 +177,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, PyList, "list", PyList::len, PyList::get_slice);

/// Used by `PyList::iter()`.
pub struct PyListIterator<'a> {
Expand Down Expand Up @@ -544,4 +532,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();
})
}
}
83 changes: 74 additions & 9 deletions src/types/sequence.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +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;
use std::ops::{Index, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};

/// Represents a reference to a Python object supporting the sequence protocol.
#[repr(transparent)]
Expand Down Expand Up @@ -254,16 +254,19 @@ impl PySequence {
}
}

impl Index<usize> for PySequence {
type Output = PyAny;
#[inline]
fn sequence_len(seq: &PySequence) -> usize {
seq.len().expect("failed to get sequence length")
}

fn index(&self, index: usize) -> &Self::Output {
self.get_item(index).unwrap_or_else(|_| {
panic!("index {} out of range for sequence", index);
})
}
#[inline]
fn sequence_slice(seq: &PySequence, start: usize, end: usize) -> &PyAny {
seq.get_slice(start, end)
.expect("sequence slice operation failed")
}

index_impls!(PySequence, PyAny, "sequence", sequence_len, sequence_slice);

impl<'a, T> FromPyObject<'a> for Vec<T>
where
T: FromPyObject<'a>,
Expand Down Expand Up @@ -439,7 +442,7 @@ mod tests {
}

#[test]
#[should_panic]
#[should_panic = "index 7 out of range for sequence"]
fn test_seq_index_trait_panic() {
Python::with_gil(|py| {
let v: Vec<i32> = vec![1, 1, 2];
Expand All @@ -449,6 +452,68 @@ mod tests {
});
}

#[test]
fn test_seq_index_trait_ranges() {
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!(vec![1, 2], seq[1..3].extract::<Vec<i32>>().unwrap());
assert_eq!(Vec::<i32>::new(), seq[3..3].extract::<Vec<i32>>().unwrap());
assert_eq!(vec![1, 2], seq[1..].extract::<Vec<i32>>().unwrap());
assert_eq!(Vec::<i32>::new(), seq[3..].extract::<Vec<i32>>().unwrap());
assert_eq!(vec![1, 1, 2], seq[..].extract::<Vec<i32>>().unwrap());
assert_eq!(vec![1, 2], seq[1..=2].extract::<Vec<i32>>().unwrap());
assert_eq!(vec![1, 1], seq[..2].extract::<Vec<i32>>().unwrap());
assert_eq!(vec![1, 1], seq[..=1].extract::<Vec<i32>>().unwrap());
})
}

#[test]
#[should_panic = "range start index 5 out of range for sequence of length 3"]
fn test_seq_index_trait_range_panic_start() {
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();
seq[5..10].extract::<Vec<i32>>().unwrap();
})
}

#[test]
#[should_panic = "range end index 10 out of range for sequence of length 3"]
fn test_seq_index_trait_range_panic_end() {
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();
seq[1..10].extract::<Vec<i32>>().unwrap();
})
}

#[test]
#[should_panic = "slice index starts at 2 but ends at 1"]
fn test_seq_index_trait_range_panic_wrong_order() {
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();
#[allow(clippy::reversed_empty_ranges)]
seq[2..1].extract::<Vec<i32>>().unwrap();
})
}

#[test]
#[should_panic = "range start index 8 out of range for sequence of length 3"]
fn test_seq_index_trait_range_from_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();
seq[8..].extract::<Vec<i32>>().unwrap();
})
}

#[test]
fn test_seq_del_item() {
Python::with_gil(|py| {
Expand Down
Loading

0 comments on commit 21759b9

Please sign in to comment.