diff --git a/tests/snippets/ints.py b/tests/snippets/ints.py index be41c3257..b947797ac 100644 --- a/tests/snippets/ints.py +++ b/tests/snippets/ints.py @@ -79,6 +79,13 @@ assert int.from_bytes(b'\x00\x10', 'little') == 4096 assert int.from_bytes(b'\xfc\x00', 'big', signed=True) == -1024 assert int.from_bytes(b'\xfc\x00', 'big', signed=False) == 64512 +assert (1024).to_bytes(4, 'big') == b'\x00\x00\x04\x00' +assert (1024).to_bytes(2, 'little', signed=True) == b'\x00\x04' +assert (-1024).to_bytes(4, 'big', signed=True) == b'\xff\xff\xfc\x00' +assert (-1024).to_bytes(4, 'little', signed=True) == b'\x00\xfc\xff\xff' +assert (2147483647).to_bytes(8, 'big', signed=False) == b'\x00\x00\x00\x00\x7f\xff\xff\xff' +assert (-2147483648).to_bytes(8, 'little', signed=True) == b'\x00\x00\x00\x80\xff\xff\xff\xff' + with assertRaises(TypeError): int(base=2) diff --git a/vm/src/obj/objint.rs b/vm/src/obj/objint.rs index 9c1ccf921..f8af3553f 100644 --- a/vm/src/obj/objint.rs +++ b/vm/src/obj/objint.rs @@ -14,6 +14,7 @@ use crate::pyobject::{ use crate::vm::VirtualMachine; use super::objbyteinner::PyByteInner; +use super::objbytes::PyBytes; use super::objstr::{PyString, PyStringRef}; use super::objtype; use crate::obj::objtype::PyClassRef; @@ -527,7 +528,74 @@ impl PyInt { } Ok(x) } + #[pymethod] + fn to_bytes( + &self, + length: PyIntRef, + byteorder: PyStringRef, + kwargs: KwArgs, + vm: &VirtualMachine, + ) -> PyResult { + let mut signed = false; + let value = self.as_bigint(); + for (key, value) in kwargs.into_iter() { + if key == "signed" { + signed = match_class!(value, + b @ PyInt => !b.as_bigint().is_zero(), + _ => false, + ); + } + } + if value.sign() == Sign::Minus && signed == false { + 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.value.as_str() { + "big" => match signed { + true => value.to_signed_bytes_be(), + false => value.to_bytes_be().1, + }, + "little" => match signed { + true => value.to_signed_bytes_le(), + false => value.to_bytes_le().1, + }, + _ => { + return Err( + vm.new_value_error("byteorder must be either 'little' or 'big'".to_string()) + ); + } + }; + let origin_len = origin_bytes.len(); + if origin_len > byte_len { + return Err(vm.new_value_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.value.as_str() { + "big" => { + bytes = append_bytes; + bytes.append(&mut origin_bytes); + } + "little" => { + bytes = origin_bytes; + bytes.append(&mut append_bytes); + } + _ => (), + } + + Ok(PyBytes::new(bytes)) + } #[pyproperty] fn real(zelf: PyRef, _vm: &VirtualMachine) -> PyIntRef { zelf