diff --git a/vm/src/obj/objint.rs b/vm/src/obj/objint.rs index e19c4f9c7..0c58ccca2 100644 --- a/vm/src/obj/objint.rs +++ b/vm/src/obj/objint.rs @@ -6,7 +6,7 @@ use num_integer::Integer; use num_traits::{Num, One, Pow, Signed, ToPrimitive, Zero}; use crate::format::FormatSpec; -use crate::function::{KwArgs, OptionalArg, PyFuncArgs}; +use crate::function::{OptionalArg, PyFuncArgs}; use crate::obj::objtype::PyClassRef; use crate::pyhash; use crate::pyobject::{ @@ -603,36 +603,30 @@ impl PyInt { }; Ok(x) } + #[pymethod] #[allow(clippy::match_bool)] - fn to_bytes( - &self, - length: PyIntRef, - byteorder: PyStringRef, - kwargs: KwArgs, - vm: &VirtualMachine, - ) -> PyResult { - let mut signed = false; + fn to_bytes(&self, options: IntToByteOptions, vm: &VirtualMachine) -> PyResult { + let signed = if let OptionalArg::Present(signed) = options.signed { + signed.to_bool() + } else { + false + }; + let value = self.as_bigint(); - for (key, value) in kwargs.into_iter() { - if key == "signed" { - signed = match_class!(match value { - b @ PyInt => !b.as_bigint().is_zero(), - _ => false, - }); - } - } if value.sign() == Sign::Minus && !signed { return Err(vm.new_overflow_error("can't convert negative int to unsigned".to_string())); } - let byte_len; - if let Some(temp) = length.as_bigint().to_usize() { - byte_len = temp; - } else { - return Err(vm.new_value_error("length parameter is illegal".to_string())); - } - let mut origin_bytes = match byteorder.as_str() { + let byte_len = if let Some(byte_len) = options.length.as_bigint().to_usize() { + byte_len + } else { + return Err( + vm.new_overflow_error("Python int too large to convert to C ssize_t".to_string()) + ); + }; + + let mut origin_bytes = match options.byteorder.as_str() { "big" => match signed { true => value.to_signed_bytes_be(), false => value.to_bytes_be().1, @@ -647,17 +641,19 @@ impl PyInt { ); } }; + let origin_len = origin_bytes.len(); if origin_len > byte_len { - return Err(vm.new_value_error("int too big to convert".to_string())); + return Err(vm.new_overflow_error("int too big to convert".to_string())); } let mut append_bytes = match value.sign() { Sign::Minus => vec![255u8; byte_len - origin_len], _ => vec![0u8; byte_len - origin_len], }; + let mut bytes = vec![]; - match byteorder.as_str() { + match options.byteorder.as_str() { "big" => { bytes = append_bytes; bytes.append(&mut origin_bytes); @@ -668,7 +664,6 @@ impl PyInt { } _ => (), } - Ok(PyBytes::new(bytes)) } #[pyproperty] @@ -738,6 +733,16 @@ struct IntFromByteOptions { signed: OptionalArg, } +#[derive(FromArgs)] +struct IntToByteOptions { + #[pyarg(positional_or_keyword)] + length: PyIntRef, + #[pyarg(positional_or_keyword)] + byteorder: PyStringRef, + #[pyarg(keyword_only, optional = true)] + signed: OptionalArg, +} + // Casting function: pub fn to_int(vm: &VirtualMachine, obj: &PyObjectRef, base: &BigInt) -> PyResult { let base_u32 = match base.to_u32() {