diff --git a/crates/vm/src/buffer.rs b/crates/vm/src/buffer.rs index 4f23e16c5..e9c557ad1 100644 --- a/crates/vm/src/buffer.rs +++ b/crates/vm/src/buffer.rs @@ -1,6 +1,6 @@ use crate::{ AsObject, Py, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine, - builtins::{PyBaseExceptionRef, PyBytesRef, PyTuple, PyTupleRef, PyType, PyTypeRef}, + builtins::{PyBaseExceptionRef, PyBytesRef, PyComplex, PyTuple, PyTupleRef, PyType, PyTypeRef}, common::{static_cell, str::wchar_t}, convert::ToPyObject, exceptions, @@ -780,7 +780,7 @@ struct PackFloatComplex(f32, f32); struct PackDoubleComplex(f64, f64); macro_rules! make_pack_complex { - ($T:ty, $Elem:ty, $Bits:ty) => { + ($T:ty, $Elem:ty, $Bits:ty, $fmt:literal) => { impl Packable for $T { fn pack( vm: &VirtualMachine, @@ -788,20 +788,37 @@ macro_rules! make_pack_complex { arg: PyObjectRef, data: &mut [u8], ) -> Result<(), PackError> { - let c = ArgIntoComplex::try_from_object(vm, arg) - .map_err(|_| { - PackError::new( - PackErrorKind::Type, - vm, - "required argument is not a complex", - ) - })? - .into_complex(); - let half = size_of::<$Elem>(); - (c.re as $Elem).to_bits().pack_int::(&mut data[..half]); - (c.im as $Elem) - .to_bits() - .pack_int::(&mut data[half..half * 2]); + let c = if let Some(value) = arg.downcast_ref::() { + value.to_complex() + } else { + ArgIntoComplex::try_from_object(vm, arg) + .map_err(|_| { + PackError::new( + PackErrorKind::Type, + vm, + "required argument is not a complex", + ) + })? + .into_complex() + }; + for (component, bytes) in [c.re, c.im] + .into_iter() + .zip(data.chunks_exact_mut(size_of::<$Elem>())) + { + let narrowed = component as $Elem; + // CPython uses native casts for matching-endian complex formats. + if E::convert(1u16) != 1 && narrowed.is_infinite() != component.is_infinite() { + return Err(PackError { + kind: PackErrorKind::Value, + exception: vm.new_overflow_error(concat!( + "float too large to pack with ", + $fmt, + " format" + )), + }); + } + narrowed.to_bits().pack_int::(bytes); + } Ok(()) } @@ -815,8 +832,8 @@ macro_rules! make_pack_complex { }; } -make_pack_complex!(PackFloatComplex, f32, u32); -make_pack_complex!(PackDoubleComplex, f64, u64); +make_pack_complex!(PackFloatComplex, f32, u32, "f"); +make_pack_complex!(PackDoubleComplex, f64, u64, "d"); impl Packable for f16 { fn pack( diff --git a/crates/vm/src/builtins/complex.rs b/crates/vm/src/builtins/complex.rs index 3f551a787..86c72d5f6 100644 --- a/crates/vm/src/builtins/complex.rs +++ b/crates/vm/src/builtins/complex.rs @@ -94,7 +94,9 @@ impl PyObjectRef { let result = method?.call((), vm)?; let ret_class = result.class().to_owned(); - if let Some(ret) = result.downcast_ref::() { + if let Some(ret) = result.downcast_ref::() + && !result.class().is(vm.ctx.types.complex_type) + { _warnings::warn( vm.ctx.exceptions.deprecation_warning, format!( diff --git a/extra_tests/snippets/stdlib_struct_complex.py b/extra_tests/snippets/stdlib_struct_complex.py new file mode 100644 index 000000000..b8ce52aca --- /dev/null +++ b/extra_tests/snippets/stdlib_struct_complex.py @@ -0,0 +1,30 @@ +import struct +import sys +import warnings + +from testutils import assert_raises + + +class ComplexLike: + def __complex__(self): + return 1 + 2j + + +class ComplexSubclass(complex): + def __complex__(self): + raise AssertionError("complex subclasses use their stored value") + + +with warnings.catch_warnings(): + warnings.simplefilter("error") + for code in ("F", "D"): + expected = struct.pack(code, 1 + 2j) + assert struct.pack(code, ComplexLike()) == expected + assert struct.pack(code, ComplexSubclass(1, 2)) == expected + +native = "<" if sys.byteorder == "little" else ">" +swapped = ">" if sys.byteorder == "little" else "<" +for value in (complex(1e300, 0), complex(0, 1e300)): + assert struct.pack(native + "F", value) == struct.pack("@F", value) + with assert_raises(OverflowError): + struct.pack(swapped + "F", value)