diff --git a/tests/snippets/builtin_range.py b/tests/snippets/builtin_range.py index 755a5e227..6a20c6f1e 100644 --- a/tests/snippets/builtin_range.py +++ b/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)) diff --git a/vm/src/obj/objrange.rs b/vm/src/obj/objrange.rs index abc6e84ab..5d5cd8de0 100644 --- a/vm/src/obj/objrange.rs +++ b/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)) } }