Implement .indices(len) of slice (Fixes #1431)

range.__getitem now also uses slice.indices() internally.
CPython: https://github.com/python/cpython/blob/master/Objects/sliceobject.c#L373
This commit is contained in:
Oscar Shrimpton
2019-10-10 15:02:09 +01:00
parent 40bbb6be1b
commit b86e803fec
2 changed files with 112 additions and 71 deletions

View File

@@ -339,77 +339,19 @@ impl PyRange {
fn getitem(&self, subscript: RangeIndex, vm: &VirtualMachine) -> PyResult {
match subscript {
RangeIndex::Slice(slice) => {
let range_start = self.start.as_bigint();
let range_step = self.step.as_bigint();
let range_length = &self.length();
let (mut substart, mut substop, mut substep) =
slice.inner_indices(&self.length(), vm)?;
let range_step = self.step(vm);
let range_start = self.start(vm);
let substep = if let Some(slice_step) = slice.step_index(vm)? {
if slice_step.is_zero() {
return Err(vm.new_value_error("slice step cannot be zero".to_string()));
}
slice_step
} else {
BigInt::one()
};
let negative_step = substep.is_negative();
let lower_bound = if negative_step {
-BigInt::one()
} else {
BigInt::zero()
};
let upper_bound = if negative_step {
&lower_bound + range_length
} else {
range_length.clone()
};
let substart = if let Some(slice_start) = slice.start_index(vm)? {
if slice_start.is_negative() {
let tmp = slice_start + range_length;
if tmp < lower_bound {
lower_bound.clone()
} else {
tmp.clone()
}
} else if slice_start > upper_bound {
upper_bound.clone()
} else {
slice_start.clone()
}
} else if negative_step {
upper_bound.clone()
} else {
lower_bound.clone()
};
let substop = if let Some(slice_stop) = slice.stop_index(vm)? {
if slice_stop.is_negative() {
let tmp = slice_stop + range_length;
if tmp < lower_bound {
lower_bound.clone()
} else {
tmp.clone()
}
} else if slice_stop > upper_bound {
upper_bound.clone()
} else {
slice_stop.clone()
}
} else if negative_step {
lower_bound.clone()
} else {
upper_bound.clone()
};
let step = range_step * &substep;
let start = range_start + (&substart * range_step);
let stop = range_start + (&substop * range_step);
substep *= range_step.as_bigint();
substart = (substart * range_step.as_bigint()) + range_start.as_bigint();
substop = (substop * range_step.as_bigint()) + range_start.as_bigint();
Ok(PyRange {
start: PyInt::new(start).into_ref(vm),
stop: PyInt::new(stop).into_ref(vm),
step: PyInt::new(step).into_ref(vm),
start: PyInt::new(substart).into_ref(vm),
stop: PyInt::new(substop).into_ref(vm),
step: PyInt::new(substep).into_ref(vm),
}
.into_ref(vm)
.into_object())

View File

@@ -1,12 +1,13 @@
use num_bigint::BigInt;
use super::objint::PyInt;
use super::objtype::{class_has_attr, PyClassRef};
use crate::function::{OptionalArg, PyFuncArgs};
use crate::pyobject::{
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef,
TypeProtocol,
};
use crate::vm::VirtualMachine;
use num_bigint::{BigInt, ToBigInt};
use num_traits::{One, Signed, Zero};
#[pyclass]
#[derive(Debug)]
@@ -157,6 +158,92 @@ impl PySlice {
Ok(eq)
}
pub(crate) fn inner_indices(
&self,
length: &BigInt,
vm: &VirtualMachine,
) -> PyResult<(BigInt, BigInt, BigInt)> {
// Calculate step
let step: BigInt;
if vm.is_none(&self.step(vm)) {
step = One::one();
} else {
// Clone the value, not the reference.
let this_step: PyRef<PyInt> = self.step(vm).try_into_ref(vm)?;
step = this_step.as_bigint().clone();
if step.is_zero() {
return Err(vm.new_value_error("slice step cannot be zero.".to_owned()));
}
}
// For convenience
let backwards = step.is_negative();
// Each end of the array
let lower = if backwards {
-1_i8.to_bigint().unwrap()
} else {
Zero::zero()
};
let upper = if backwards {
lower.clone() + length
} else {
length.clone()
};
// Calculate start
let mut start: BigInt;
if vm.is_none(&self.start(vm)) {
// Default
start = if backwards {
upper.clone()
} else {
lower.clone()
};
} else {
let this_start: PyRef<PyInt> = self.start(vm).try_into_ref(vm)?;
start = this_start.as_bigint().clone();
if start < Zero::zero() {
// From end of array
start += length;
if start < lower {
start = lower.clone();
}
} else if start > upper {
start = upper.clone();
}
}
// Calculate Stop
let mut stop: BigInt;
if vm.is_none(&self.stop(vm)) {
stop = if backwards {
lower.clone()
} else {
upper.clone()
};
} else {
let this_stop: PyRef<PyInt> = self.stop(vm).try_into_ref(vm)?;
stop = this_stop.as_bigint().clone();
if stop < Zero::zero() {
// From end of array
stop += length;
if stop < lower {
stop = lower.clone();
}
} else if stop > upper {
stop = upper.clone();
}
}
Ok((start, stop, step))
}
#[pymethod(name = "__eq__")]
fn eq(&self, rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult {
if let Some(rhs) = rhs.payload::<PySlice>() {
@@ -221,6 +308,18 @@ impl PySlice {
fn hash(&self, vm: &VirtualMachine) -> PyResult<()> {
Err(vm.new_type_error("unhashable type".to_string()))
}
#[pymethod(name = "indices")]
fn indices(&self, length: PyObjectRef, vm: &VirtualMachine) -> PyResult {
if let Some(length) = length.payload::<PyInt>() {
let (start, stop, step) = self.inner_indices(length.as_bigint(), vm)?;
Ok(vm
.ctx
.new_tuple(vec![vm.new_int(start), vm.new_int(stop), vm.new_int(step)]))
} else {
Ok(vm.ctx.not_implemented())
}
}
}
fn to_index_value(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Option<BigInt>> {