Merge pull request RustPython#353 from OddCoincidence/range-index

Add range.index

Author: windelbouwman

Cargo.lock

@@ -1 +1,31 @@
+def assert_raises(expr, exc_type):
+    """
+    Helper function to assert `expr` raises an exception of type `exc_type`
+    Args:
+        expr: Callable
+        exec_type: Exception
+    Returns:
+        None
+    Raises:
+        Assertion error on failure
+    """
+    try:
+        expr(None)
+    except exc_type:
+        assert True
+    else:
+        assert False
+
 assert range(2**63+1)[2**63] == 9223372036854775808
+
+# index tests
+assert range(10).index(6) == 6
+assert range(4, 10).index(6) == 2
+assert range(4, 10, 2).index(6) == 1
+
+# index raises value error on out of bounds
+assert_raises(lambda _: range(10).index(-1), ValueError)
+assert_raises(lambda _: range(10).index(10), ValueError)
+
+# index raises value error if out of step
+assert_raises(lambda _: range(4, 10, 2).index(5), ValueError)

tests/snippets/builtin_range.py

@@ -8,6 +8,7 @@ bitflags = "1.0.4"
 num-complex = "0.2"
 num-bigint = "0.2.1"
 num-traits = "0.2"
+num-integer = "0.1.39"
 rand = "0.5"
 log = "0.3"
 rustpython_parser = {path = "../parser"}

vm/Cargo.toml

@@ -12,6 +12,7 @@ extern crate log;
 // extern crate env_logger;
 extern crate num_bigint;
 extern crate num_complex;
+extern crate num_integer;
 extern crate num_traits;
 extern crate serde;
 extern crate serde_json;

vm/src/lib.rs

@@ -5,6 +5,7 @@ use super::super::vm::VirtualMachine;
 use super::objint;
 use super::objtype;
 use num_bigint::{BigInt, ToBigInt};
+use num_integer::Integer;
 use num_traits::{One, Signed, ToPrimitive, Zero};
 
 #[derive(Debug, Clone)]
@@ -25,6 +26,20 @@ impl RangeType {
             .unwrap()
     }
 
+    #[inline]
+    pub fn index_of(&self, value: &BigInt) -> Option<BigInt> {
+        if value < &self.start || value >= &self.end {
+            return None;
+        }
+
+        let offset = value - &self.start;
+        if offset.is_multiple_of(&self.step) {
+            Some(offset / &self.step)
+        } else {
+            None
+        }
+    }
+
     #[inline]
     pub fn is_empty(&self) -> bool {
         (self.start <= self.end && self.step.is_negative())
@@ -60,6 +75,7 @@ pub fn init(context: &PyContext) {
         "__getitem__",
         context.new_rustfunc(range_getitem),
     );
+    context.set_attr(&range_type, "index", context.new_rustfunc(range_index));
 }
 
 fn range_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -193,3 +209,23 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         _ => Err(vm.new_type_error("range indices must be integer or slice".to_string())),
     }
 }
+
+fn range_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.range_type())), (needle, None)]
+    );
+
+    if let PyObjectPayload::Range { ref range } = zelf.borrow().payload {
+        match needle.borrow().payload {
+            PyObjectPayload::Integer { ref value } => match range.index_of(value) {
+                Some(idx) => Ok(vm.ctx.new_int(idx)),
+                None => Err(vm.new_value_error(format!("{} is not in range", value))),
+            },
+            _ => Err(vm.new_value_error("sequence.index(x): x not in sequence".to_string())),
+        }
+    } else {
+        unreachable!()
+    }
+}