Merge pull request #381 from nficca/implement-range-count

Implement range.count

tests/snippets/builtin_range.py

@@ -39,6 +39,15 @@ assert_raises(lambda _: range(4, 10, 2).index(5), ValueError)
 # index raises value error if needle is not an int
 assert_raises(lambda _: range(10).index('foo'), ValueError)
 
+# count tests
+assert range(10).count(2) == 1
+assert range(10).count(11) == 0
+assert range(10).count(-1) == 0
+assert range(9, 12).count(10) == 1
+assert range(4, 10, 2).count(4) == 1
+assert range(4, 10, 2).count(7) == 0
+assert range(10).count("foo") == 0
+
 # __bool__
 assert bool(range(1))
 assert bool(range(1, 2))

vm/src/obj/objrange.rs

@@ -66,6 +66,14 @@ impl RangeType {
         }
     }
 
+    #[inline]
+    pub fn count(&self, value: &BigInt) -> usize {
+        match self.index_of(value).is_some() {
+            true => 1,
+            false => 0,
+        }
+    }
+
     #[inline]
     pub fn is_empty(&self) -> bool {
         (self.start <= self.end && self.step.is_negative())
@@ -128,6 +136,14 @@ impl RangeType {
     }
 }
 
+pub fn get_value(obj: &PyObjectRef) -> RangeType {
+    if let PyObjectPayload::Range { range } = &obj.borrow().payload {
+        range.clone()
+    } else {
+        panic!("Inner error getting range {:?}", obj);
+    }
+}
+
 pub fn init(context: &PyContext) {
     let ref range_type = context.range_type;
 
@@ -165,6 +181,7 @@ pub fn init(context: &PyContext) {
         context.new_rustfunc(range_contains),
     );
     context.set_attr(&range_type, "index", context.new_rustfunc(range_index));
+    context.set_attr(&range_type, "count", context.new_rustfunc(range_count));
 }
 
 fn range_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -223,10 +240,7 @@ fn range_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn range_reversed(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zelf, Some(vm.ctx.range_type()))]);
 
-    let range = match zelf.borrow().payload {
-        PyObjectPayload::Range { ref range } => range.reversed(),
-        _ => unreachable!(),
-    };
+    let range = get_value(zelf).reversed();
 
     Ok(PyObject::new(
         PyObjectPayload::Iterator {
@@ -240,10 +254,7 @@ fn range_reversed(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn range_len(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zelf, Some(vm.ctx.range_type()))]);
 
-    if let Some(len) = match zelf.borrow().payload {
-        PyObjectPayload::Range { ref range } => range.try_len(),
-        _ => unreachable!(),
-    } {
+    if let Some(len) = get_value(zelf).try_len() {
         Ok(vm.ctx.new_int(len))
     } else {
         Err(vm.new_overflow_error("Python int too large to convert to Rust usize".to_string()))
@@ -256,15 +267,12 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         args,
         required = [(zelf, Some(vm.ctx.range_type())), (subscript, None)]
     );
-    let zrange = if let PyObjectPayload::Range { ref range } = zelf.borrow().payload {
-        range.clone()
-    } else {
-        unreachable!()
-    };
+
+    let range = get_value(zelf);
 
     match subscript.borrow().payload {
         PyObjectPayload::Integer { ref value } => {
-            if let Some(int) = zrange.get(value) {
+            if let Some(int) = range.get(value) {
                 Ok(vm.ctx.new_int(int))
             } else {
                 Err(vm.new_index_error("range object index out of range".to_string()))
@@ -276,29 +284,29 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
             ref step,
         } => {
             let new_start = if let Some(int) = start {
-                if let Some(i) = zrange.get(int) {
+                if let Some(i) = range.get(int) {
                     i
                 } else {
-                    zrange.start.clone()
+                    range.start.clone()
                 }
             } else {
-                zrange.start.clone()
+                range.start.clone()
             };
 
             let new_end = if let Some(int) = stop {
-                if let Some(i) = zrange.get(int) {
+                if let Some(i) = range.get(int) {
                     i
                 } else {
-                    zrange.end
+                    range.end
                 }
             } else {
-                zrange.end
+                range.end
             };
 
             let new_step = if let Some(int) = step {
-                int * zrange.step
+                int * range.step
             } else {
-                zrange.step
+                range.step
             };
 
             Ok(PyObject::new(
@@ -320,21 +328,15 @@ fn range_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn range_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zelf, Some(vm.ctx.range_type()))]);
 
-    let s = match zelf.borrow().payload {
-        PyObjectPayload::Range { ref range } => range.repr(),
-        _ => unreachable!(),
-    };
+    let repr = get_value(zelf).repr();
 
-    Ok(vm.ctx.new_str(s))
+    Ok(vm.ctx.new_str(repr))
 }
 
 fn range_bool(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zelf, Some(vm.ctx.range_type()))]);
 
-    let len = match zelf.borrow().payload {
-        PyObjectPayload::Range { ref range } => range.len(),
-        _ => unreachable!(),
-    };
+    let len = get_value(zelf).len();
 
     Ok(vm.ctx.new_bool(len > 0))
 }
@@ -346,14 +348,15 @@ fn range_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(zelf, Some(vm.ctx.range_type())), (needle, None)]
     );
 
-    if let PyObjectPayload::Range { ref range } = zelf.borrow().payload {
-        Ok(vm.ctx.new_bool(match needle.borrow().payload {
-            PyObjectPayload::Integer { ref value } => range.contains(value),
-            _ => false,
-        }))
+    let range = get_value(zelf);
+
+    let result = if objtype::isinstance(needle, &vm.ctx.int_type()) {
+        range.contains(&objint::get_value(needle))
     } else {
-        unreachable!()
-    }
+        false
+    };
+
+    Ok(vm.ctx.new_bool(result))
 }
 
 fn range_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -363,15 +366,32 @@ fn range_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         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())),
+    let range = get_value(zelf);
+
+    if objtype::isinstance(needle, &vm.ctx.int_type()) {
+        let needle = objint::get_value(needle);
+
+        match range.index_of(&needle) {
+            Some(idx) => Ok(vm.ctx.new_int(idx)),
+            None => Err(vm.new_value_error(format!("{} is not in range", needle))),
         }
     } else {
-        unreachable!()
+        Err(vm.new_value_error("sequence.index(x): x not in sequence".to_string()))
+    }
+}
+
+fn range_count(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(zelf, Some(vm.ctx.range_type())), (item, None)]
+    );
+
+    let range = get_value(zelf);
+
+    if objtype::isinstance(item, &vm.ctx.int_type()) {
+        Ok(vm.ctx.new_int(range.count(&objint::get_value(item))))
+    } else {
+        Ok(vm.ctx.new_int(0))
     }
 }