Function spec_format_bytes

crates/vm/src/cformat.rs:21–124

(
    vm: &VirtualMachine,
    spec: &CFormatSpec,
    obj: PyObjectRef,
)

use num_traits::cast::ToPrimitive;

fn spec_format_bytes(
    vm: &VirtualMachine,
    spec: &CFormatSpec,
    obj: PyObjectRef,
) -> PyResult<Vec<u8>> {
    match &spec.format_type {
        CFormatType::String(conversion) => match conversion {
            // Unlike strings, %r and %a are identical for bytes: the behaviour corresponds to
            // %a for strings (not %r)
            CFormatConversion::Repr | CFormatConversion::Ascii => {
                let b = builtins::ascii(obj, vm)?.as_bytes().to_vec();
                Ok(b)
            }
            CFormatConversion::Str | CFormatConversion::Bytes => {
                if let Ok(buffer) = PyBuffer::try_from_borrowed_object(vm, &obj) {
                    Ok(buffer.contiguous_or_collect(|bytes| spec.format_bytes(bytes)))
                } else {
                    let bytes = vm
                        .get_special_method(&obj, identifier!(vm, __bytes__))?
                        .ok_or_else(|| {
                            vm.new_type_error(format!(
                                "%b requires a bytes-like object, or an object that \
                                    implements __bytes__, not '{}'",
                                obj.class().name()
                            ))
                        })?
                        .invoke((), vm)?;
                    let bytes = PyBytes::try_from_borrowed_object(vm, &bytes)?;
                    Ok(spec.format_bytes(bytes.as_bytes()))
                }
            }
        },
        CFormatType::Number(number_type) => match number_type {
            CNumberType::DecimalD | CNumberType::DecimalI | CNumberType::DecimalU => {
                match_class!(match &obj {
                    ref i @ PyInt => {
                        Ok(spec.format_number(i.as_bigint()).into_bytes())
                    }
                    ref f @ PyFloat => {
                        Ok(spec
                            .format_number(&try_f64_to_bigint(f.to_f64(), vm)?)
                            .into_bytes())
                    }
                    obj => {
                        if let Some(method) = vm.get_method(obj.clone(), identifier!(vm, __int__)) {
                            let result = method?.call((), vm)?;
                            if let Some(i) = result.downcast_ref::<PyInt>() {
                                return Ok(spec.format_number(i.as_bigint()).into_bytes());
                            }
                        }
                        Err(vm.new_type_error(format!(
                            "%{} format: a number is required, not {}",
                            spec.format_type.to_char(),
                            obj.class().name()
                        )))
                    }
                })
            }