rustpython_vm/
byte.rs

1//! byte operation APIs
2
3use num_traits::ToPrimitive;
4
5use crate::{
6    AsObject, PyObject, PyObjectRef, PyResult, VirtualMachine,
7    protocol::{BufferFlags, PyBuffer},
8};
9
10/// The element error `bytes` reports, which is the one entry point here that does
11/// not name a single byte: `bytes([256])` says "bytes must be in range(0, 256)"
12/// where every `bytearray` path says "byte".
13const BYTES_ELEMENT_ERROR: &str = "bytes must be in range(0, 256)";
14const BYTE_ELEMENT_ERROR: &str = "byte must be in range(0, 256)";
15
16// PyBytes_FromObject
17pub fn bytes_from_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Vec<u8>> {
18    collect_bytes(vm, obj, true, BYTES_ELEMENT_ERROR, |name| {
19        format!("cannot convert '{name}' object to bytes")
20    })
21}
22
23/// [`bytes_from_object`] for the bytearray constructor and for assigning to a
24/// slice of one, which run the iterator without asking the object they were
25/// handed how long it is.
26pub fn bytearray_from_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Vec<u8>> {
27    collect_bytes(vm, obj, false, BYTE_ELEMENT_ERROR, |name| {
28        format!("cannot convert '{name}' object to bytearray")
29    })
30}
31
32/// [`bytes_from_object`] for `bytearray_extend()`, which names what it was
33/// doing rather than what it was converting to.
34pub fn bytearray_extend_from_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Vec<u8>> {
35    collect_bytes(vm, obj, true, BYTE_ELEMENT_ERROR, |name| {
36        format!("can't extend bytearray with {name}")
37    })
38}
39
40/// `measured` is whether the object is asked how long it is; `element` is the error
41/// for a value outside `range(0, 256)`; `unusable` names, from the class name, what
42/// could not be done with one that is not iterable.
43fn collect_bytes(
44    vm: &VirtualMachine,
45    obj: &PyObject,
46    measured: bool,
47    element: &'static str,
48    unusable: impl FnOnce(&str) -> String,
49) -> PyResult<Vec<u8>> {
50    if obj.check_buffer() {
51        let buffer = PyBuffer::from_object(vm, obj, BufferFlags::FULL_RO)?;
52        return Ok(buffer.contiguous_or_collect(|bytes| bytes.to_vec()));
53    }
54
55    if !obj.fast_isinstance(vm.ctx.types.str_type) {
56        // What `PyObject_GetIter()` cannot take is answered for by the caller,
57        // which knows what it was being asked to do, rather than by the
58        // iteration protocol saying the object is not iterable.
59        let cls = obj.class();
60        if cls.slots().iter.load().is_none() && !cls.has_attr(identifier!(vm, __getitem__)) {
61            return Err(vm.new_type_error(unusable(&cls.name())));
62        }
63        let value = |x: PyObjectRef| value_from_object_with(vm, &x, element);
64        let elements = if measured {
65            vm.map_iterable_object_sized(obj, value)
66        } else {
67            vm.map_iterable_object(obj, value)
68        };
69        if let Ok(elements) = elements {
70            return elements;
71        }
72    }
73
74    Err(vm.new_type_error("can assign only bytes, buffers, or iterables of ints in range(0, 256)"))
75}
76
77pub fn value_from_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<u8> {
78    value_from_object_with(vm, obj, BYTE_ELEMENT_ERROR)
79}
80
81fn value_from_object_with(
82    vm: &VirtualMachine,
83    obj: &PyObject,
84    element: &'static str,
85) -> PyResult<u8> {
86    obj.try_index(vm)?
87        .as_bigint()
88        .to_u8()
89        .ok_or_else(|| vm.new_value_error(element))
90}
rustpython_vm/byte.rs

rustpython_vm/
byte.rs
