forked from Rust-related/RustPython
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:
@@ -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())
|
||||
|
||||
@@ -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>> {
|
||||
|
||||
Reference in New Issue
Block a user