Skip to content

Commit c01a867

Browse files
Merge pull request RustPython#974 from stwen77/add_int_to_bytes
Add to_bytes for int
2 parents 002140a + 68a1a9f commit c01a867

2 files changed

Lines changed: 75 additions & 0 deletions

File tree

tests/snippets/ints.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,13 @@
7979
assert int.from_bytes(b'\xfc\x00', 'big', signed=True) == -1024
8080
assert int.from_bytes(b'\xfc\x00', 'big', signed=False) == 64512
8181

82+
assert (1024).to_bytes(4, 'big') == b'\x00\x00\x04\x00'
83+
assert (1024).to_bytes(2, 'little', signed=True) == b'\x00\x04'
84+
assert (-1024).to_bytes(4, 'big', signed=True) == b'\xff\xff\xfc\x00'
85+
assert (-1024).to_bytes(4, 'little', signed=True) == b'\x00\xfc\xff\xff'
86+
assert (2147483647).to_bytes(8, 'big', signed=False) == b'\x00\x00\x00\x00\x7f\xff\xff\xff'
87+
assert (-2147483648).to_bytes(8, 'little', signed=True) == b'\x00\x00\x00\x80\xff\xff\xff\xff'
88+
8289
with assertRaises(TypeError):
8390
int(base=2)
8491

vm/src/obj/objint.rs

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use crate::pyobject::{
1414
use crate::vm::VirtualMachine;
1515

1616
use super::objbyteinner::PyByteInner;
17+
use super::objbytes::PyBytes;
1718
use super::objstr::{PyString, PyStringRef};
1819
use super::objtype;
1920
use crate::obj::objtype::PyClassRef;
@@ -527,7 +528,74 @@ impl PyInt {
527528
}
528529
Ok(x)
529530
}
531+
#[pymethod]
532+
fn to_bytes(
533+
&self,
534+
length: PyIntRef,
535+
byteorder: PyStringRef,
536+
kwargs: KwArgs,
537+
vm: &VirtualMachine,
538+
) -> PyResult<PyBytes> {
539+
let mut signed = false;
540+
let value = self.as_bigint();
541+
for (key, value) in kwargs.into_iter() {
542+
if key == "signed" {
543+
signed = match_class!(value,
544+
545+
b @ PyInt => !b.as_bigint().is_zero(),
546+
_ => false,
547+
);
548+
}
549+
}
550+
if value.sign() == Sign::Minus && signed == false {
551+
return Err(vm.new_overflow_error("can't convert negative int to unsigned".to_string()));
552+
}
553+
let byte_len;
554+
if let Some(temp) = length.as_bigint().to_usize() {
555+
byte_len = temp;
556+
} else {
557+
return Err(vm.new_value_error("length parameter is illegal".to_string()));
558+
}
559+
560+
let mut origin_bytes = match byteorder.value.as_str() {
561+
"big" => match signed {
562+
true => value.to_signed_bytes_be(),
563+
false => value.to_bytes_be().1,
564+
},
565+
"little" => match signed {
566+
true => value.to_signed_bytes_le(),
567+
false => value.to_bytes_le().1,
568+
},
569+
_ => {
570+
return Err(
571+
vm.new_value_error("byteorder must be either 'little' or 'big'".to_string())
572+
);
573+
}
574+
};
575+
let origin_len = origin_bytes.len();
576+
if origin_len > byte_len {
577+
return Err(vm.new_value_error("int too big to convert".to_string()));
578+
}
579+
580+
let mut append_bytes = match value.sign() {
581+
Sign::Minus => vec![255u8; byte_len - origin_len],
582+
_ => vec![0u8; byte_len - origin_len],
583+
};
584+
let mut bytes = vec![];
585+
match byteorder.value.as_str() {
586+
"big" => {
587+
bytes = append_bytes;
588+
bytes.append(&mut origin_bytes);
589+
}
590+
"little" => {
591+
bytes = origin_bytes;
592+
bytes.append(&mut append_bytes);
593+
}
594+
_ => (),
595+
}
530596

597+
Ok(PyBytes::new(bytes))
598+
}
531599
#[pyproperty]
532600
fn real(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyIntRef {
533601
zelf

0 commit comments

Comments
 (0)