diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs index 5e21b04e7..d3dd0e732 100644 --- a/vm/src/stdlib/mod.rs +++ b/vm/src/stdlib/mod.rs @@ -16,6 +16,7 @@ mod json; mod keyword; mod marshal; mod math; +mod operator; mod platform; mod pystruct; mod random; @@ -74,6 +75,7 @@ pub fn get_module_inits() -> HashMap { "json".to_string() => Box::new(json::make_module), "marshal".to_string() => Box::new(marshal::make_module), "math".to_string() => Box::new(math::make_module), + "_operator".to_string() => Box::new(operator::make_module), "platform".to_string() => Box::new(platform::make_module), "regex_crate".to_string() => Box::new(re::make_module), "random".to_string() => Box::new(random::make_module), diff --git a/vm/src/stdlib/operator.rs b/vm/src/stdlib/operator.rs new file mode 100644 index 000000000..b2fbc7bea --- /dev/null +++ b/vm/src/stdlib/operator.rs @@ -0,0 +1,24 @@ +use crate::function::OptionalArg; +use crate::obj::{objiter, objtype}; +use crate::pyobject::{PyObjectRef, PyResult, TypeProtocol}; +use crate::VirtualMachine; + +fn operator_length_hint(obj: PyObjectRef, default: OptionalArg, vm: &VirtualMachine) -> PyResult { + let default = default.unwrap_or_else(|| vm.new_int(0)); + if !objtype::isinstance(&default, &vm.ctx.types.int_type) { + return Err(vm.new_type_error(format!( + "'{}' type cannot be interpreted as an integer", + default.class().name + ))); + } + let hint = objiter::length_hint(vm, obj)? + .map(|i| vm.new_int(i)) + .unwrap_or(default); + Ok(hint) +} + +pub fn make_module(vm: &VirtualMachine) -> PyObjectRef { + py_module!(vm, "_operator", { + "length_hint" => vm.ctx.new_rustfunc(operator_length_hint), + }) +}