diff --git a/crates/stdlib/src/math.rs b/crates/stdlib/src/math.rs index 9ea0be522c6..cb50261a3b1 100644 --- a/crates/stdlib/src/math.rs +++ b/crates/stdlib/src/math.rs @@ -6,13 +6,11 @@ use crate::vm::{VirtualMachine, builtins::PyBaseExceptionRef}; #[pymodule] mod math { use crate::vm::{ - AsObject, Py, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine, + AsObject, Py, PyObject, PyObjectRef, PyResult, VirtualMachine, builtins::{ PyFloat, PyInt, PyIntRef, PyModule, PyStrInterned, try_bigint_to_f64, try_f64_to_bigint, }, - function::{ - ArgIndex, ArgIntoFloat, ArgIterable, Either, NameCoordinates, OptionalArg, PosArgs, - }, + function::{ArgIndex, ArgIntoFloat, ArgIterable, NameCoordinates, OptionalArg, PosArgs}, identifier, }; use malachite_bigint::BigInt; @@ -457,7 +455,7 @@ mod math { #[derive(FromArgs)] struct LdexpArgs { #[pyarg(positional)] - x: Either, PyIntRef>, + x: ArgIntoFloat, // Refuses anything that is not an `int`, including objects with `__index__`. #[pyarg(positional, error_msg = "Expected an int as second argument to ldexp.")] i: PyIntRef, @@ -466,11 +464,8 @@ mod math { #[pyfunction] fn ldexp(args: LdexpArgs, vm: &VirtualMachine) -> PyResult { let LdexpArgs { x, i } = args; - let value = match x { - Either::A(f) => f.to_f64(), - Either::B(z) => try_bigint_to_f64(z.as_bigint(), vm)?, - }; - pymath::math::ldexp_bigint(value, i.as_bigint()).map_err(|err| pymath_exception(err, vm)) + pymath::math::ldexp_bigint(x.into_float(), i.as_bigint()) + .map_err(|err| pymath_exception(err, vm)) } #[pyfunction] diff --git a/crates/vm/src/function/number.rs b/crates/vm/src/function/number.rs index 9fdb06a9717..53203a683d3 100644 --- a/crates/vm/src/function/number.rs +++ b/crates/vm/src/function/number.rs @@ -1,7 +1,7 @@ use super::argument::OptionalArg; use crate::{ AsObject, Py, PyObjectRef, PyResult, TryFromObject, VirtualMachine, - builtins::{PyInt, PyIntRef}, + builtins::{PyFloat, PyInt, PyIntRef}, }; use core::ops::Deref; use malachite_bigint::BigInt; @@ -88,6 +88,12 @@ impl From for f64 { impl TryFromObject for ArgIntoFloat { // Equivalent to PyFloat_AsDouble. fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { + // A float subclass is read directly, without calling its __float__. + if let Some(float) = obj.downcast_ref::() { + return Ok(Self { + value: float.to_f64(), + }); + } let value = obj.try_float(vm)?.to_f64(); Ok(Self { value }) } diff --git a/extra_tests/snippets/stdlib_math.py b/extra_tests/snippets/stdlib_math.py index b06bf8446f6..7b41f039e83 100644 --- a/extra_tests/snippets/stdlib_math.py +++ b/extra_tests/snippets/stdlib_math.py @@ -541,3 +541,33 @@ def assertAllNotClose(examples, *args, **kwargs): assert math.perm(5) == math.perm(5, None) == 120 assert math.perm(2**64, 0) == 1 assert math.comb(2**64, 2**64 - 1) == 2**64 + + +class LdexpFloat: + def __float__(self): + return 1.5 + + +class LdexpIndex: + def __index__(self): + return 3 + + +class LdexpRaisingFloat: + def __float__(self): + raise RuntimeError + + +assert math.ldexp(LdexpFloat(), 1) == 3.0 +assert math.ldexp(LdexpIndex(), 1) == 6.0 +assert_raises(RuntimeError, math.ldexp, LdexpRaisingFloat(), 1) +assert_raises(OverflowError, math.ldexp, 10**400, 1) + + +class FloatSubclass(float): + def __float__(self): + return 9.0 + + +assert math.sqrt(FloatSubclass(4.0)) == 2.0 +assert math.ldexp(FloatSubclass(1.0), 1) == 2.0