[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/RustPython/RustPython-test/main/crates/stdlib/src/array.rs [Back]  [Original]

// spell-checker:ignore typecode tofile tolist fromfile

use rustpython_vm::{PyRef, VirtualMachine, builtins::PyModule};

pub(crate) fn make_module(vm: &VirtualMachine) -> PyRef {
    let module = array::make_module(vm);

    let array = module
        .get_attr("array", vm)
        .expect("Expect array has array type.");

    let collections_abc = vm
        .import("collections.abc", 0)
        .expect("Expect collections exist.");
    let abc = collections_abc
        .get_attr("abc", vm)
        .expect("Expect collections has abc submodule.");
    let mutable_sequence = abc
        .get_attr("MutableSequence", vm)
        .expect("Expect collections.abc has MutableSequence type.");

    let register = &mutable_sequence
        .get_attr("register", vm)
        .expect("Expect collections.abc.MutableSequence has register method.");
    register
        .call((array,), vm)
        .expect("Expect collections.abc.MutableSequence.register(array.array) not fail.");

    module
}

#[pymodule(name = "array")]
mod array {
    use crate::{
        common::{
            atomic::{self, AtomicUsize},
            lock::{
                PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,
                PyRwLockReadGuard, PyRwLockWriteGuard,
            },
            str::wchar_t,
        },
        vm::{
            AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
            atomic_func,
            builtins::{
                PositionIterInternal, PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat,
                PyGenericAlias, PyInt, PyList, PyListRef, PyStr, PyStrRef, PyTupleRef, PyTypeRef,
            },
            class_or_notimplemented,
            convert::{ToPyObject, ToPyResult, TryFromBorrowedObject, TryFromObject},
            function::{
                ArgBytesLike, ArgIntoFloat, ArgIterable, KwArgs, OptionalArg, PyComparisonValue,
            },
            protocol::{
                BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
                PyMappingMethods, PySequenceMethods,
            },
            sequence::{OptionalRangeArgs, SequenceExt, SequenceMutExt},
            sliceable::{
                SaturatedSlice, SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp,
                SliceableSequenceOp,
            },
            types::{
                AsBuffer, AsMapping, AsSequence, Comparable, Constructor, IterNext, Iterable,
                PyComparisonOp, Representable, SelfIter,
            },
        },
    };
    use itertools::Itertools;
    use num_traits::ToPrimitive;
    use rustpython_common::wtf8::{CodePoint, Wtf8, Wtf8Buf};
    use std::{cmp::Ordering, fmt, os::raw};

    macro_rules! def_array_enum {
        ($(($n:ident, $t:ty, $c:literal, $scode:literal)),*$(,)?) => {
            #[derive(Debug, Clone)]
            pub enum ArrayContentType {
                $($n(Vec),)*
            }

            impl ArrayContentType {
                fn from_char(c: char) -> Result {
                    match c {
                        $($c => Ok(ArrayContentType::$n(Vec::new())),)*
                        _ => Err(
                            "bad typecode (must be b, B, u, h, H, i, I, l, L, q, Q, f or d)".into()
                        ),
                    }
                }

                const fn typecode(&self) -> char {
                    match self {
                        $(ArrayContentType::$n(_) => $c,)*
                    }
                }

                const fn typecode_str(&self) -> &'static str {
                    match self {
                        $(ArrayContentType::$n(_) => $scode,)*
                    }
                }

                const fn itemsize_of_typecode(c: char) -> Option {
                    match c {
                        $($c => Some(std::mem::size_of::()),)*
                        _ => None,
                    }
                }

                const fn itemsize(&self) -> usize {
                    match self {
                        $(ArrayContentType::$n(_) => std::mem::size_of::(),)*
                    }
                }

                fn addr(&self) -> usize {
                    match self {
                        $(ArrayContentType::$n(v) => v.as_ptr() as usize,)*
                    }
                }

                fn len(&self) -> usize {
                    match self {
                        $(ArrayContentType::$n(v) => v.len(),)*
                    }
                }

                fn reserve(&mut self, len: usize) {
                    match self {
                        $(ArrayContentType::$n(v) => v.reserve(len),)*
                    }
                }

                fn push(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            let val = ::try_into_from_object(vm, obj)?;
                            v.push(val);
                        })*
                    }
                    Ok(())
                }

                fn pop(&mut self, i: isize, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                        let i = v.wrap_index(i).ok_or_else(|| {
                            vm.new_index_error("pop index out of range".to_owned())
                        })?;
                            v.remove(i).to_pyresult(vm)
                        })*
                    }
                }

                fn insert(
                    &mut self,
                    i: isize,
                    obj: PyObjectRef,
                    vm: &VirtualMachine
                ) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            let val = ::try_into_from_object(vm, obj)?;
                            v.insert(i.saturated_at(v.len()), val);
                        })*
                    }
                    Ok(())
                }

                fn count(&self, obj: PyObjectRef, vm: &VirtualMachine) -> usize {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            if let Ok(val) = ::try_into_from_object(vm, obj) {
                                v.iter().filter(|&&a| a == val).count()
                            } else {
                                0
                            }
                        })*
                    }
                }

                fn remove(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult{
                    match self {
                        $(ArrayContentType::$n(v) => {
                            if let Ok(val) = ::try_into_from_object(vm, obj) {
                                if let Some(pos) = v.iter().position(|&a| a == val) {
                                    v.remove(pos);
                                    return Ok(());
                                }
                            }
                            Err(vm.new_value_error("array.remove(x): x not in array".to_owned()))
                        })*
                    }
                }

                fn frombytes_move(&mut self, b: Vec) {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            if v.is_empty() {
                                // safe because every configuration of bytes for the types we
                                // support are valid
                                let b = std::mem::ManuallyDrop::new(b);
                                let ptr = b.as_ptr() as *mut $t;
                                let len = b.len() / std::mem::size_of::();
                                let capacity = b.capacity() / std::mem::size_of::();
                                *v = unsafe { Vec::from_raw_parts(ptr, len, capacity) };
                            } else {
                                self.frombytes(&b);
                            }
                        })*
                    }
                }

                fn frombytes(&mut self, b: &[u8]) {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // safe because every configuration of bytes for the types we
                            // support are valid
                            if b.len() > 0 {
                                let ptr = b.as_ptr() as *const $t;
                                let ptr_len = b.len() / std::mem::size_of::();
                                let slice = unsafe { std::slice::from_raw_parts(ptr, ptr_len) };
                                v.extend_from_slice(slice);
                            }
                        })*
                    }
                }

                fn fromlist(&mut self, list: &PyList, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // convert list before modify self
                            let mut list: Vec = list
                                .borrow_vec()
                                .iter()
                                .cloned()
                                .map(|value| ::try_into_from_object(vm, value))
                                .try_collect()?;
                            v.append(&mut list);
                            Ok(())
                        })*
                    }
                }

                fn get_bytes(&self) -> &[u8] {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // safe because we're just reading memory as bytes
                            let ptr = v.as_ptr() as *const u8;
                            let ptr_len = v.len() * std::mem::size_of::();
                            unsafe { std::slice::from_raw_parts(ptr, ptr_len) }
                        })*
                    }
                }

                fn get_bytes_mut(&mut self) -> &mut [u8] {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // safe because we're just reading memory as bytes
                            let ptr = v.as_ptr() as *mut u8;
                            let ptr_len = v.len() * std::mem::size_of::();
                            unsafe { std::slice::from_raw_parts_mut(ptr, ptr_len) }
                        })*
                    }
                }

                fn index(
                    &self,
                    obj: PyObjectRef,
                    start: usize,
                    stop: usize,
                    vm: &VirtualMachine
                ) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            if let Ok(val) = ::try_into_from_object(vm, obj) {
                                if let Some(pos) = v.iter().take(stop as _).skip(start as _).position(|&elem| elem == val) {
                                    return Ok(pos + start);
                                }
                            }
                            Err(vm.new_value_error("array.index(x): x not in array".to_owned()))
                        })*
                    }
                }

                fn reverse(&mut self) {
                    match self {
                        $(ArrayContentType::$n(v) => v.reverse(),)*
                    }
                }

                fn get(
                    &self,
                    i: usize,
                    vm: &VirtualMachine
                ) -> Option {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            v.get(i).map(|x| x.to_pyresult(vm))
                        })*
                    }
                }

                fn getitem_by_index(&self, i: isize, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            v.getitem_by_index(vm, i).map(|x| x.to_pyresult(vm))?
                        })*
                    }
                }

                fn getitem_by_slice(&self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            let r = v.getitem_by_slice(vm, slice)?;
                            let array = PyArray::from(ArrayContentType::$n(r));
                            array.to_pyresult(vm)
                        })*
                    }
                }

                fn setitem_by_index(
                    &mut self,
                    i: isize,
                    value: PyObjectRef,
                    vm: &VirtualMachine
                ) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            let value = ::try_into_from_object(vm, value)?;
                            v.setitem_by_index(vm, i, value)
                        })*
                    }
                }

                fn setitem_by_slice(
                    &mut self,
                    slice: SaturatedSlice,
                    items: &ArrayContentType,
                    vm: &VirtualMachine
                ) -> PyResult {
                    match self {
                        $(Self::$n(elements) => if let ArrayContentType::$n(items) = items {
                            elements.setitem_by_slice(vm, slice, items)
                        } else {
                            Err(vm.new_type_error(
                                "bad argument type for built-in operation".to_owned()
                            ))
                        },)*
                    }
                }

                fn setitem_by_slice_no_resize(
                    &mut self,
                    slice: SaturatedSlice,
                    items: &ArrayContentType,
                    vm: &VirtualMachine
                ) -> PyResult {
                    match self {
                        $(Self::$n(elements) => if let ArrayContentType::$n(items) = items {
                            elements.setitem_by_slice_no_resize(vm, slice, items)
                        } else {
                            Err(vm.new_type_error(
                                "bad argument type for built-in operation".to_owned()
                            ))
                        },)*
                    }
                }

                fn delitem_by_index(&mut self, i: isize, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            v.delitem_by_index(vm, i)
                        })*
                    }
                }

                fn delitem_by_slice(&mut self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            v.delitem_by_slice(vm, slice)
                        })*
                    }
                }

                fn add(&self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {
                            let elements = v.iter().chain(other.iter()).cloned().collect();
                            Ok(ArrayContentType::$n(elements))
                        } else {
                            Err(vm.new_type_error(
                                "bad argument type for built-in operation".to_owned()
                            ))
                        },)*
                    }
                }

                fn iadd(&mut self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {
                            v.extend(other);
                            Ok(())
                        } else {
                            Err(vm.new_type_error(
                                "can only extend with array of same kind".to_owned()
                            ))
                        },)*
                    }
                }

                fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // MemoryError instead Overflow Error, hard to says it is right
                            // but it is how cpython doing right now
                            let elements = v.mul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))?;
                            Ok(ArrayContentType::$n(elements))
                        })*
                    }
                }

                fn imul(&mut self, value: isize, vm: &VirtualMachine) -> PyResult {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            // MemoryError instead Overflow Error, hard to says it is right
                            // but it is how cpython doing right now
                            v.imul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))
                        })*
                    }
                }

                fn byteswap(&mut self) {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            for element in v.iter_mut() {
                                let x = element.byteswap();
                                *element = x;
                            }
                        })*
                    }
                }

                fn repr(&self, class_name: &str, _vm: &VirtualMachine) -> PyResult {
                    // we don't need ReprGuard here
                    let s = match self {
                        $(ArrayContentType::$n(v) => {
                            if v.is_empty() {
                                format!("{}('{}')", class_name, $c)
                            } else {
                                format!("{}('{}', [{}])", class_name, $c, v.iter().format(", "))
                            }
                        })*
                    };
                    Ok(s)
                }

                fn iter Result {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            if let ArrayContentType::$n(other) = other {
                                Ok(PartialOrd::partial_cmp(v, other))
                            } else {
                                Err(())
                            }
                        })*
                    }
                }

                fn get_objects(&self, vm: &VirtualMachine) -> Vec {
                    match self {
                        $(ArrayContentType::$n(v) => {
                            v.iter().map(|&x| x.to_object(vm)).collect()
                        })*
                    }
                }
            }
        };
    }

    def_array_enum!(
        (SignedByte, i8, 'b', "b"),
        (UnsignedByte, u8, 'B', "B"),
        (PyUnicode, WideChar, 'u', "u"),
        (SignedShort, raw::c_short, 'h', "h"),
        (UnsignedShort, raw::c_ushort, 'H', "H"),
        (SignedInt, raw::c_int, 'i', "i"),
        (UnsignedInt, raw::c_uint, 'I', "I"),
        (SignedLong, raw::c_long, 'l', "l"),
        (UnsignedLong, raw::c_ulong, 'L', "L"),
        (SignedLongLong, raw::c_longlong, 'q', "q"),
        (UnsignedLongLong, raw::c_ulonglong, 'Q', "Q"),
        (Float, f32, 'f', "f"),
        (Double, f64, 'd', "d"),
    );

    #[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
    pub struct WideChar(wchar_t);

    trait ArrayElement: Sized {
        fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult;
        fn byteswap(self) -> Self;
        fn to_object(self, vm: &VirtualMachine) -> PyObjectRef;
    }

    macro_rules! impl_int_element {
        ($($t:ty,)*) => {$(
            impl ArrayElement for $t {
                fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult {
                    obj.try_index(vm)?.try_to_primitive(vm)
                }
                fn byteswap(self) -> Self {
                    ::swap_bytes(self)
                }
                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
                    self.to_pyobject(vm)
                }
            }
        )*};
    }

    macro_rules! impl_float_element {
        ($(($t:ty, $f_from:path, $f_swap:path, $f_to:path),)*) => {$(
            impl ArrayElement for $t {
                fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult {
                    $f_from(vm, obj)
                }
                fn byteswap(self) -> Self {
                    $f_swap(self)
                }
                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
                    $f_to(self).into_pyobject(vm)
                }
            }
        )*};
    }

    impl_int_element!(i8, u8, i16, u16, i32, u32, i64, u64,);
    impl_float_element!(
        (
            f32,
            f32_try_into_from_object,
            f32_swap_bytes,
            pyfloat_from_f32
        ),
        (f64, f64_try_into_from_object, f64_swap_bytes, PyFloat::from),
    );

    const fn f32_swap_bytes(x: f32) -> f32 {
        f32::from_bits(x.to_bits().swap_bytes())
    }

    const fn f64_swap_bytes(x: f64) -> f64 {
        f64::from_bits(x.to_bits().swap_bytes())
    }

    fn f32_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult {
        ArgIntoFloat::try_from_object(vm, obj).map(|x| *x as f32)
    }

    fn f64_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult {
        ArgIntoFloat::try_from_object(vm, obj).map(Into::into)
    }

    fn pyfloat_from_f32(value: f32) -> PyFloat {
        PyFloat::from(value as f64)
    }

    impl ArrayElement for WideChar {
        fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult {
            PyStrRef::try_from_object(vm, obj)?
                .as_str()
                .chars()
                .exactly_one()
                .map(|ch| Self(ch as _))
                .map_err(|_| vm.new_type_error("array item must be unicode character"))
        }
        fn byteswap(self) -> Self {
            Self(self.0.swap_bytes())
        }
        fn to_object(self, _vm: &VirtualMachine) -> PyObjectRef {
            unreachable!()
        }
    }

    fn u32_to_char(ch: u32) -> Result {
        CodePoint::from_u32(ch)
            .ok_or_else(|| format!("character U+{ch:4x} is not in range [U+0000; U+10ffff]"))
    }

    impl TryFrom for CodePoint {
        type Error = String;

        fn try_from(ch: WideChar) -> Result {
            // safe because every configuration of bytes for the types we support are valid
            u32_to_char(ch.0 as _)
        }
    }

    impl ToPyResult for WideChar {
        fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {
            Ok(CodePoint::try_from(self)
                .map_err(|e| vm.new_unicode_encode_error(e))?
                .to_pyobject(vm))
        }
    }

    impl fmt::Display for WideChar {
        fn fmt(&self, _f: &mut fmt::Formatter {
            self.array.read()
        }

        fn write(&self) -> PyRwLockWriteGuard {
            PyRwLockReadGuard::map(self.read(), |a| a.get_bytes())
        }

        pub(crate) fn get_bytes_mut(&self) -> PyMappedRwLockWriteGuard &'static PySequenceMethods {
            static AS_SEQUENCE: PySequenceMethods = PySequenceMethods {
                length: atomic_func!(|seq, _vm| Ok(PyArray::sequence_downcast(seq).__len__())),
                concat: atomic_func!(|seq, other, vm| {
                    let zelf = PyArray::sequence_downcast(seq);
                    PyArray::__add__(zelf, other.to_owned(), vm).map(|x| x.into())
                }),
                repeat: atomic_func!(|seq, n, vm| {
                    PyArray::sequence_downcast(seq)
                        .__mul__(n, vm)
                        .map(|x| x.into())
                }),
                item: atomic_func!(|seq, i, vm| {
                    PyArray::sequence_downcast(seq)
                        .read()
                        .getitem_by_index(i, vm)
                }),
                ass_item: atomic_func!(|seq, i, value, vm| {
                    let zelf = PyArray::sequence_downcast(seq);
                    if let Some(value) = value {
                        zelf.write().setitem_by_index(i, value, vm)
                    } else {
                        zelf.write().delitem_by_index(i, vm)
                    }
                }),
                contains: atomic_func!(|seq, target, vm| {
                    let zelf = PyArray::sequence_downcast(seq);
                    Ok(zelf.__contains__(target.to_owned(), vm))
                }),
                inplace_concat: atomic_func!(|seq, other, vm| {
                    let zelf = PyArray::sequence_downcast(seq).to_owned();
                    PyArray::__iadd__(zelf, other.to_owned(), vm).map(|x| x.into())
                }),
                inplace_repeat: atomic_func!(|seq, n, vm| {
                    let zelf = PyArray::sequence_downcast(seq).to_owned();
                    PyArray::__imul__(zelf, n, vm).map(|x| x.into())
                }),
            };
            &AS_SEQUENCE
        }
    }

    impl Iterable for PyArray {
        fn iter(zelf: PyRef, vm: &VirtualMachine) -> PyResult {
            Ok(PyArrayIter {
                internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
            }
            .into_pyobject(vm))
        }
    }

    impl BufferResizeGuard for PyArray {
        type Resizable;

        fn try_resizable_opt(&self) -> Option TryFromBorrowedObject PyResult {
            obj.try_to_ref::(vm)
                .map_err(|_| {
                    vm.new_type_error(format!(
                        "an integer is required (got type {})",
                        obj.class().name()
                    ))
                })?
                .try_to_primitive::(vm)?
                .to_u8()
                .unwrap_or(u8::MAX)
                .try_into()
                .map_err(|_| {
                    vm.new_value_error("third argument must be a valid machine format code.")
                })
        }
    }

    impl MachineFormatCode {
        fn from_typecode(code: char) -> Option {
            use std::mem::size_of;
            let signed = code.is_ascii_uppercase();
            let big_endian = cfg!(target_endian = "big");
            let int_size = match code {
                'b' | 'B' => return Some(Self::Int8 { signed }),
                'u' => {
                    return match size_of::() {
                        2 => Some(Self::Utf16 { big_endian }),
                        4 => Some(Self::Utf32 { big_endian }),
                        _ => None,
                    };
                }
                'f' => {
                    // Copied from CPython
                    const Y: f32 = 16711938.0;
                    return match &Y.to_ne_bytes() {
                        b"\x4b\x7f\x01\x02" => Some(Self::Ieee754Float { big_endian: true }),
                        b"\x02\x01\x7f\x4b" => Some(Self::Ieee754Float { big_endian: false }),
                        _ => None,
                    };
                }
                'd' => {
                    // Copied from CPython
                    const Y: f64 = 9006104071832581.0;
                    return match &Y.to_ne_bytes() {
                        b"\x43\x3f\xff\x01\x02\x03\x04\x05" => {
                            Some(Self::Ieee754Double { big_endian: true })
                        }
                        b"\x05\x04\x03\x02\x01\xff\x3f\x43" => {
                            Some(Self::Ieee754Double { big_endian: false })
                        }
                        _ => None,
                    };
                }
                _ => ArrayContentType::itemsize_of_typecode(code)? as u8,
            };
            match int_size {
                2 => Some(Self::Int16 { signed, big_endian }),
                4 => Some(Self::Int32 { signed, big_endian }),
                8 => Some(Self::Int64 { signed, big_endian }),
                _ => None,
            }
        }
        const fn item_size(self) -> usize {
            match self {
                Self::Int8 { .. } => 1,
                Self::Int16 { .. } | Self::Utf16 { .. } => 2,
                Self::Int32 { .. } | Self::Utf32 { .. } | Self::Ieee754Float { .. } => 4,
                Self::Int64 { .. } | Self::Ieee754Double { .. } => 8,
            }
        }
    }

    fn check_array_type(typ: PyTypeRef, vm: &VirtualMachine) -> PyResult {
        if !typ.fast_issubclass(PyArray::class(&vm.ctx)) {
            return Err(
                vm.new_type_error(format!("{} is not a subtype of array.array", typ.name()))
            );
        }
        Ok(typ)
    }

    fn check_type_code(spec: PyStrRef, vm: &VirtualMachine) -> PyResult {
        let spec = spec.as_str().chars().exactly_one().map_err(|_| {
            vm.new_type_error(
                "_array_reconstructor() argument 2 must be a unicode character, not str",
            )
        })?;
        ArrayContentType::from_char(spec)
            .map_err(|_| vm.new_value_error("second argument must be a valid type code"))
    }

    macro_rules! chunk_to_obj {
        ($BYTE:ident, $TY:ty, $BIG_ENDIAN:ident) => {{
            let b = ::try_from($BYTE).unwrap();
            if $BIG_ENDIAN {
                ::from_be_bytes(b)
            } else {
                ::from_le_bytes(b)
            }
        }};
        ($VM:ident, $BYTE:ident, $TY:ty, $BIG_ENDIAN:ident) => {
            chunk_to_obj!($BYTE, $TY, $BIG_ENDIAN).to_pyobject($VM)
        };
        ($VM:ident, $BYTE:ident, $SIGNED_TY:ty, $UNSIGNED_TY:ty, $SIGNED:ident, $BIG_ENDIAN:ident) => {{
            let b = ::try_from($BYTE).unwrap();
            match ($SIGNED, $BIG_ENDIAN) {
                (false, false) => ::from_le_bytes(b).to_pyobject($VM),
                (false, true) => ::from_be_bytes(b).to_pyobject($VM),
                (true, false) => ::from_le_bytes(b).to_pyobject($VM),
                (true, true) => ::from_be_bytes(b).to_pyobject($VM),
            }
        }};
    }

    #[pyfunction]
    fn _array_reconstructor(args: ReconstructorArgs, vm: &VirtualMachine) -> PyResult {
        let cls = check_array_type(args.arraytype, vm)?;
        let mut array = check_type_code(args.typecode, vm)?;
        let format = args.mformat_code;
        let bytes = args.items.as_bytes();
        if !bytes.len().is_multiple_of(format.item_size()) {
            return Err(vm.new_value_error("bytes length not a multiple of item size"));
        }
        if MachineFormatCode::from_typecode(array.typecode()) == Some(format) {
            array.frombytes(bytes);
            return PyArray::from(array).into_ref_with_type(vm, cls);
        }
        if !matches!(
            format,
            MachineFormatCode::Utf16 { .. } | MachineFormatCode::Utf32 { .. }
        ) {
            array.reserve(bytes.len() / format.item_size());
        }
        let mut chunks = bytes.chunks(format.item_size());
        match format {
            MachineFormatCode::Ieee754Float { big_endian } => {
                chunks.try_for_each(|b| array.push(chunk_to_obj!(vm, b, f32, big_endian), vm))?
            }
            MachineFormatCode::Ieee754Double { big_endian } => {
                chunks.try_for_each(|b| array.push(chunk_to_obj!(vm, b, f64, big_endian), vm))?
            }
            MachineFormatCode::Int8 { signed } => chunks
                .try_for_each(|b| array.push(chunk_to_obj!(vm, b, i8, u8, signed, false), vm))?,
            MachineFormatCode::Int16 { signed, big_endian } => chunks.try_for_each(|b| {
                array.push(chunk_to_obj!(vm, b, i16, u16, signed, big_endian), vm)
            })?,
            MachineFormatCode::Int32 { signed, big_endian } => chunks.try_for_each(|b| {
                array.push(chunk_to_obj!(vm, b, i32, u32, signed, big_endian), vm)
            })?,
            MachineFormatCode::Int64 { signed, big_endian } => chunks.try_for_each(|b| {
                array.push(chunk_to_obj!(vm, b, i64, u64, signed, big_endian), vm)
            })?,
            MachineFormatCode::Utf16 { big_endian } => {
                let utf16: Vec = chunks.map(|b| chunk_to_obj!(b, u16, big_endian)).collect();
                let s = String::from_utf16(&utf16)
                    .map_err(|_| vm.new_unicode_encode_error("items cannot decode as utf16"))?;
                let bytes = PyArray::_unicode_to_wchar_bytes((*s).as_ref(), array.itemsize());
                array.frombytes_move(bytes);
            }
            MachineFormatCode::Utf32 { big_endian } => {
                let s: Wtf8Buf = chunks
                    .map(|b| chunk_to_obj!(b, u32, big_endian))
                    .map(|ch| u32_to_char(ch).map_err(|msg| vm.new_value_error(msg)))
                    .try_collect()?;
                let bytes = PyArray::_unicode_to_wchar_bytes(&s, array.itemsize());
                array.frombytes_move(bytes);
            }
        };
        PyArray::from(array).into_ref_with_type(vm, cls)
    }
}

Web Proxy Viewer  |  New URL  |  Original Page