Skip to main content

rustpython_stdlib/
array.rs

1// spell-checker:ignore typecode tofile tolist fromfile
2
3pub(crate) use array::module_def;
4
5#[pymodule(name = "array")]
6mod array {
7    use crate::{
8        common::{
9            atomic::{self, AtomicUsize},
10            lock::{
11                PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,
12                PyRwLockReadGuard, PyRwLockWriteGuard,
13            },
14            str::wchar_t,
15        },
16        vm::{
17            AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
18            atomic_func,
19            builtins::{
20                PositionIterInternal, PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat,
21                PyGenericAlias, PyInt, PyList, PyListRef, PyStr, PyStrRef, PyTupleRef, PyType,
22                PyTypeRef, PyUtf8StrRef, builtins_iter,
23            },
24            class_or_notimplemented,
25            convert::{ToPyObject, ToPyResult, TryFromBorrowedObject, TryFromObject},
26            function::{
27                ArgBytesLike, ArgIntoFloat, ArgIterable, KwArgs, OptionalArg, PyComparisonValue,
28            },
29            protocol::{
30                BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
31                PyMappingMethods, PySequenceMethods,
32            },
33            sequence::{OptionalRangeArgs, SequenceExt, SequenceMutExt},
34            sliceable::{
35                SaturatedSlice, SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp,
36                SliceableSequenceOp,
37            },
38            stdlib::_warnings,
39            types::{
40                AsBuffer, AsMapping, AsSequence, Comparable, Constructor, IterNext, Iterable,
41                PyComparisonOp, Representable, SelfIter,
42            },
43        },
44    };
45    use alloc::fmt;
46    use core::cmp::Ordering;
47    use itertools::Itertools;
48    use num_traits::ToPrimitive;
49    use rustpython_common::wtf8::{CodePoint, Wtf8, Wtf8Buf};
50    use std::os::raw;
51    macro_rules! def_array_enum {
52        ($(($n:ident, $t:ty, $c:literal, $scode:literal)),*$(,)?) => {
53            #[derive(Debug, Clone)]
54            pub enum ArrayContentType {
55                $($n(Vec<$t>),)*
56            }
57
58            impl ArrayContentType {
59                fn from_char(c: char) -> Result<Self, String> {
60                    match c {
61                        $($c => Ok(ArrayContentType::$n(Vec::new())),)*
62                        _ => Err(
63                            "bad typecode (must be b, B, u, h, H, i, I, l, L, q, Q, f or d)".into()
64                        ),
65                    }
66                }
67
68                const fn typecode(&self) -> char {
69                    match self {
70                        $(ArrayContentType::$n(_) => $c,)*
71                    }
72                }
73
74                const fn typecode_str(&self) -> &'static str {
75                    match self {
76                        $(ArrayContentType::$n(_) => $scode,)*
77                    }
78                }
79
80                const fn itemsize_of_typecode(c: char) -> Option<usize> {
81                    match c {
82                        $($c => Some(core::mem::size_of::<$t>()),)*
83                        _ => None,
84                    }
85                }
86
87                const fn itemsize(&self) -> usize {
88                    match self {
89                        $(ArrayContentType::$n(_) => core::mem::size_of::<$t>(),)*
90                    }
91                }
92
93                fn addr(&self) -> usize {
94                    match self {
95                        $(ArrayContentType::$n(v) => v.as_ptr() as usize,)*
96                    }
97                }
98
99                fn len(&self) -> usize {
100                    match self {
101                        $(ArrayContentType::$n(v) => v.len(),)*
102                    }
103                }
104
105                fn reserve(&mut self, len: usize) {
106                    match self {
107                        $(ArrayContentType::$n(v) => v.reserve(len),)*
108                    }
109                }
110
111                fn push(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
112                    match self {
113                        $(ArrayContentType::$n(v) => {
114                            let val = <$t>::try_into_from_object(vm, obj)?;
115                            v.push(val);
116                        })*
117                    }
118                    Ok(())
119                }
120
121                fn pop(&mut self, i: isize, vm: &VirtualMachine) -> PyResult {
122                    match self {
123                        $(ArrayContentType::$n(v) => {
124                        let i = v.wrap_index(i).ok_or_else(|| {
125                            vm.new_index_error("pop index out of range".to_owned())
126                        })?;
127                            v.remove(i).to_pyresult(vm)
128                        })*
129                    }
130                }
131
132                fn insert(
133                    &mut self,
134                    i: isize,
135                    obj: PyObjectRef,
136                    vm: &VirtualMachine
137                ) -> PyResult<()> {
138                    match self {
139                        $(ArrayContentType::$n(v) => {
140                            let val = <$t>::try_into_from_object(vm, obj)?;
141                            v.insert(i.saturated_at(v.len()), val);
142                        })*
143                    }
144                    Ok(())
145                }
146
147                fn count(&self, obj: PyObjectRef, vm: &VirtualMachine) -> usize {
148                    match self {
149                        $(ArrayContentType::$n(v) => {
150                            if let Ok(val) = <$t>::try_into_from_object(vm, obj) {
151                                v.iter().filter(|&&a| a == val).count()
152                            } else {
153                                0
154                            }
155                        })*
156                    }
157                }
158
159                fn clear(&mut self) -> PyResult<()>{
160                    match self {
161                        $(ArrayContentType::$n(v) => v.clear(),)*
162                    };
163                    Ok(())
164                }
165
166                fn remove(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()>{
167                    match self {
168                        $(ArrayContentType::$n(v) => {
169                            if let Ok(val) = <$t>::try_into_from_object(vm, obj) {
170                                if let Some(pos) = v.iter().position(|&a| a == val) {
171                                    v.remove(pos);
172                                    return Ok(());
173                                }
174                            }
175                            Err(vm.new_value_error("array.remove(x): x not in array".to_owned()))
176                        })*
177                    }
178                }
179
180                fn frombytes_move(&mut self, b: Vec<u8>) {
181                    match self {
182                        $(ArrayContentType::$n(v) => {
183                            if v.is_empty() {
184                                // safe because every configuration of bytes for the types we
185                                // support are valid
186                                let b = core::mem::ManuallyDrop::new(b);
187                                let ptr = b.as_ptr() as *mut $t;
188                                let len = b.len() / core::mem::size_of::<$t>();
189                                let capacity = b.capacity() / core::mem::size_of::<$t>();
190                                *v = unsafe { Vec::from_raw_parts(ptr, len, capacity) };
191                            } else {
192                                self.frombytes(&b);
193                            }
194                        })*
195                    }
196                }
197
198                fn frombytes(&mut self, b: &[u8]) {
199                    match self {
200                        $(ArrayContentType::$n(v) => {
201                            // safe because every configuration of bytes for the types we
202                            // support are valid
203                            if b.len() > 0 {
204                                let ptr = b.as_ptr() as *const $t;
205                                let ptr_len = b.len() / core::mem::size_of::<$t>();
206                                let slice = unsafe { core::slice::from_raw_parts(ptr, ptr_len) };
207                                v.extend_from_slice(slice);
208                            }
209                        })*
210                    }
211                }
212
213                fn fromlist(&mut self, list: &PyList, vm: &VirtualMachine) -> PyResult<()> {
214                    match self {
215                        $(ArrayContentType::$n(v) => {
216                            // convert list before modify self
217                            let mut list: Vec<$t> = list
218                                .borrow_vec()
219                                .iter()
220                                .cloned()
221                                .map(|value| <$t>::try_into_from_object(vm, value))
222                                .try_collect()?;
223                            v.append(&mut list);
224                            Ok(())
225                        })*
226                    }
227                }
228
229                fn get_bytes(&self) -> &[u8] {
230                    match self {
231                        $(ArrayContentType::$n(v) => {
232                            // safe because we're just reading memory as bytes
233                            let ptr = v.as_ptr() as *const u8;
234                            let ptr_len = v.len() * core::mem::size_of::<$t>();
235                            unsafe { core::slice::from_raw_parts(ptr, ptr_len) }
236                        })*
237                    }
238                }
239
240                fn get_bytes_mut(&mut self) -> &mut [u8] {
241                    match self {
242                        $(ArrayContentType::$n(v) => {
243                            // safe because we're just reading memory as bytes
244                            let ptr = v.as_ptr() as *mut u8;
245                            let ptr_len = v.len() * core::mem::size_of::<$t>();
246                            unsafe { core::slice::from_raw_parts_mut(ptr, ptr_len) }
247                        })*
248                    }
249                }
250
251                fn index(
252                    &self,
253                    obj: PyObjectRef,
254                    start: usize,
255                    stop: usize,
256                    vm: &VirtualMachine
257                ) -> PyResult<usize> {
258                    match self {
259                        $(ArrayContentType::$n(v) => {
260                            if let Ok(val) = <$t>::try_into_from_object(vm, obj) {
261                                if let Some(pos) = v.iter().take(stop as _).skip(start as _).position(|&elem| elem == val) {
262                                    return Ok(pos + start);
263                                }
264                            }
265                            Err(vm.new_value_error("array.index(x): x not in array".to_owned()))
266                        })*
267                    }
268                }
269
270                fn reverse(&mut self) {
271                    match self {
272                        $(ArrayContentType::$n(v) => v.reverse(),)*
273                    }
274                }
275
276                fn get(
277                    &self,
278                    i: usize,
279                    vm: &VirtualMachine
280                ) -> Option<PyResult> {
281                    match self {
282                        $(ArrayContentType::$n(v) => {
283                            v.get(i).map(|x| x.to_pyresult(vm))
284                        })*
285                    }
286                }
287
288                fn getitem_by_index(&self, i: isize, vm: &VirtualMachine) -> PyResult {
289                    match self {
290                        $(ArrayContentType::$n(v) => {
291                            v.getitem_by_index(vm, i).map(|x| x.to_pyresult(vm))?
292                        })*
293                    }
294                }
295
296                fn getitem_by_slice(&self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult {
297                    match self {
298                        $(ArrayContentType::$n(v) => {
299                            let r = v.getitem_by_slice(vm, slice)?;
300                            let array = PyArray::from(ArrayContentType::$n(r));
301                            array.to_pyresult(vm)
302                        })*
303                    }
304                }
305
306                fn setitem_by_index(
307                    &mut self,
308                    i: isize,
309                    value: PyObjectRef,
310                    vm: &VirtualMachine
311                ) -> PyResult<()> {
312                    match self {
313                        $(ArrayContentType::$n(v) => {
314                            let value = <$t>::try_into_from_object(vm, value)?;
315                            v.setitem_by_index(vm, i, value)
316                        })*
317                    }
318                }
319
320                fn setitem_by_slice(
321                    &mut self,
322                    slice: SaturatedSlice,
323                    items: &ArrayContentType,
324                    vm: &VirtualMachine
325                ) -> PyResult<()> {
326                    match self {
327                        $(Self::$n(elements) => if let ArrayContentType::$n(items) = items {
328                            elements.setitem_by_slice(vm, slice, items)
329                        } else {
330                            Err(vm.new_type_error(
331                                "bad argument type for built-in operation".to_owned()
332                            ))
333                        },)*
334                    }
335                }
336
337                fn setitem_by_slice_no_resize(
338                    &mut self,
339                    slice: SaturatedSlice,
340                    items: &ArrayContentType,
341                    vm: &VirtualMachine
342                ) -> PyResult<()> {
343                    match self {
344                        $(Self::$n(elements) => if let ArrayContentType::$n(items) = items {
345                            elements.setitem_by_slice_no_resize(vm, slice, items)
346                        } else {
347                            Err(vm.new_type_error(
348                                "bad argument type for built-in operation".to_owned()
349                            ))
350                        },)*
351                    }
352                }
353
354                fn delitem_by_index(&mut self, i: isize, vm: &VirtualMachine) -> PyResult<()> {
355                    match self {
356                        $(ArrayContentType::$n(v) => {
357                            v.delitem_by_index(vm, i)
358                        })*
359                    }
360                }
361
362                fn delitem_by_slice(&mut self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult<()> {
363                    match self {
364                        $(ArrayContentType::$n(v) => {
365                            v.delitem_by_slice(vm, slice)
366                        })*
367                    }
368                }
369
370                fn add(&self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult<Self> {
371                    match self {
372                        $(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {
373                            let elements = v.iter().chain(other.iter()).cloned().collect();
374                            Ok(ArrayContentType::$n(elements))
375                        } else {
376                            Err(vm.new_type_error(
377                                "bad argument type for built-in operation".to_owned()
378                            ))
379                        },)*
380                    }
381                }
382
383                fn iadd(&mut self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult<()> {
384                    match self {
385                        $(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {
386                            v.extend(other);
387                            Ok(())
388                        } else {
389                            Err(vm.new_type_error(
390                                "can only extend with array of same kind".to_owned()
391                            ))
392                        },)*
393                    }
394                }
395
396                fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
397                    match self {
398                        $(ArrayContentType::$n(v) => {
399                            // MemoryError instead Overflow Error, hard to says it is right
400                            // but it is how cpython doing right now
401                            let elements = v.mul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))?;
402                            Ok(ArrayContentType::$n(elements))
403                        })*
404                    }
405                }
406
407                fn imul(&mut self, value: isize, vm: &VirtualMachine) -> PyResult<()> {
408                    match self {
409                        $(ArrayContentType::$n(v) => {
410                            // MemoryError instead Overflow Error, hard to says it is right
411                            // but it is how cpython doing right now
412                            v.imul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))
413                        })*
414                    }
415                }
416
417                fn byteswap(&mut self) {
418                    match self {
419                        $(ArrayContentType::$n(v) => {
420                            for element in v.iter_mut() {
421                                let x = element.byteswap();
422                                *element = x;
423                            }
424                        })*
425                    }
426                }
427
428                fn repr(&self, class_name: &str, _vm: &VirtualMachine) -> PyResult<String> {
429                    // we don't need ReprGuard here
430                    let s = match self {
431                        $(ArrayContentType::$n(v) => {
432                            if v.is_empty() {
433                                format!("{}('{}')", class_name, $c)
434                            } else {
435                                format!("{}('{}', [{}])", class_name, $c, v.iter().format(", "))
436                            }
437                        })*
438                    };
439                    Ok(s)
440                }
441
442                fn iter<'a, 'vm: 'a>(
443                    &'a self,
444                    vm: &'vm VirtualMachine
445                ) -> impl Iterator<Item = PyResult> + 'a {
446                    (0..self.len()).map(move |i| self.get(i, vm).unwrap())
447                }
448
449                fn cmp(&self, other: &ArrayContentType) -> Result<Option<Ordering>, ()> {
450                    match self {
451                        $(ArrayContentType::$n(v) => {
452                            if let ArrayContentType::$n(other) = other {
453                                Ok(PartialOrd::partial_cmp(v, other))
454                            } else {
455                                Err(())
456                            }
457                        })*
458                    }
459                }
460
461                fn get_objects(&self, vm: &VirtualMachine) -> Vec<PyObjectRef> {
462                    match self {
463                        $(ArrayContentType::$n(v) => {
464                            v.iter().map(|&x| x.to_object(vm)).collect()
465                        })*
466                    }
467                }
468            }
469        };
470    }
471
472    def_array_enum!(
473        (SignedByte, i8, 'b', "b"),
474        (UnsignedByte, u8, 'B', "B"),
475        (PyUnicode, WideChar, 'u', "u"),
476        (SignedShort, raw::c_short, 'h', "h"),
477        (UnsignedShort, raw::c_ushort, 'H', "H"),
478        (SignedInt, raw::c_int, 'i', "i"),
479        (UnsignedInt, raw::c_uint, 'I', "I"),
480        (SignedLong, raw::c_long, 'l', "l"),
481        (UnsignedLong, raw::c_ulong, 'L', "L"),
482        (SignedLongLong, raw::c_longlong, 'q', "q"),
483        (UnsignedLongLong, raw::c_ulonglong, 'Q', "Q"),
484        (Float, f32, 'f', "f"),
485        (Double, f64, 'd', "d"),
486    );
487
488    #[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]
489    pub struct WideChar(wchar_t);
490
491    trait ArrayElement: Sized {
492        fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self>;
493        fn byteswap(self) -> Self;
494        fn to_object(self, vm: &VirtualMachine) -> PyObjectRef;
495    }
496
497    macro_rules! impl_int_element {
498        ($($t:ty,)*) => {$(
499            impl ArrayElement for $t {
500                fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
501                    obj.try_index(vm)?.try_to_primitive(vm)
502                }
503                fn byteswap(self) -> Self {
504                    <$t>::swap_bytes(self)
505                }
506                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
507                    self.to_pyobject(vm)
508                }
509            }
510        )*};
511    }
512
513    macro_rules! impl_float_element {
514        ($(($t:ty, $f_from:path, $f_swap:path, $f_to:path),)*) => {$(
515            impl ArrayElement for $t {
516                fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
517                    $f_from(vm, obj)
518                }
519                fn byteswap(self) -> Self {
520                    $f_swap(self)
521                }
522                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
523                    $f_to(self).into_pyobject(vm)
524                }
525            }
526        )*};
527    }
528
529    impl_int_element!(i8, u8, i16, u16, i32, u32, i64, u64,);
530    impl_float_element!(
531        (
532            f32,
533            f32_try_into_from_object,
534            f32_swap_bytes,
535            pyfloat_from_f32
536        ),
537        (f64, f64_try_into_from_object, f64_swap_bytes, PyFloat::from),
538    );
539
540    const fn f32_swap_bytes(x: f32) -> f32 {
541        f32::from_bits(x.to_bits().swap_bytes())
542    }
543
544    const fn f64_swap_bytes(x: f64) -> f64 {
545        f64::from_bits(x.to_bits().swap_bytes())
546    }
547
548    fn f32_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<f32> {
549        ArgIntoFloat::try_from_object(vm, obj).map(|x| x.into_float() as f32)
550    }
551
552    fn f64_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<f64> {
553        ArgIntoFloat::try_from_object(vm, obj).map(|x| x.into_float())
554    }
555
556    fn pyfloat_from_f32(value: f32) -> PyFloat {
557        PyFloat::from(value as f64)
558    }
559
560    impl ArrayElement for WideChar {
561        fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
562            PyUtf8StrRef::try_from_object(vm, obj)?
563                .as_str()
564                .chars()
565                .exactly_one()
566                .map(|ch| Self(ch as _))
567                .map_err(|_| vm.new_type_error("array item must be unicode character"))
568        }
569        fn byteswap(self) -> Self {
570            Self(self.0.swap_bytes())
571        }
572        fn to_object(self, _vm: &VirtualMachine) -> PyObjectRef {
573            unreachable!()
574        }
575    }
576
577    fn u32_to_char(ch: u32) -> Result<CodePoint, String> {
578        CodePoint::from_u32(ch)
579            .ok_or_else(|| format!("character U+{ch:4x} is not in range [U+0000; U+10ffff]"))
580    }
581
582    impl TryFrom<WideChar> for CodePoint {
583        type Error = String;
584
585        fn try_from(ch: WideChar) -> Result<Self, Self::Error> {
586            // safe because every configuration of bytes for the types we support are valid
587            u32_to_char(ch.0 as _)
588        }
589    }
590
591    impl ToPyResult for WideChar {
592        fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {
593            Ok(CodePoint::try_from(self)
594                .map_err(|e| vm.new_unicode_encode_error(e))?
595                .to_pyobject(vm))
596        }
597    }
598
599    impl fmt::Display for WideChar {
600        fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
601            unreachable!("`repr(array('u'))` calls `PyStr::repr`")
602        }
603    }
604
605    #[pyattr]
606    #[pyattr(name = "ArrayType")]
607    #[pyclass(name = "array")]
608    #[derive(Debug, PyPayload)]
609    pub struct PyArray {
610        array: PyRwLock<ArrayContentType>,
611        exports: AtomicUsize,
612    }
613
614    pub type PyArrayRef = PyRef<PyArray>;
615
616    impl From<ArrayContentType> for PyArray {
617        fn from(array: ArrayContentType) -> Self {
618            Self {
619                array: PyRwLock::new(array),
620                exports: AtomicUsize::new(0),
621            }
622        }
623    }
624
625    #[derive(FromArgs)]
626    pub struct ArrayNewArgs {
627        #[pyarg(positional)]
628        spec: PyUtf8StrRef,
629        #[pyarg(positional, optional)]
630        init: OptionalArg<PyObjectRef>,
631    }
632
633    impl Constructor for PyArray {
634        type Args = (ArrayNewArgs, KwArgs);
635
636        fn py_new(
637            cls: &Py<PyType>,
638            (ArrayNewArgs { spec, init }, kwargs): Self::Args,
639            vm: &VirtualMachine,
640        ) -> PyResult<Self> {
641            let spec = spec.as_str().chars().exactly_one().map_err(|_| {
642                vm.new_type_error("array() argument 1 must be a unicode character, not str")
643            })?;
644
645            if cls.is(Self::class(&vm.ctx)) && !kwargs.is_empty() {
646                return Err(vm.new_type_error("array.array() takes no keyword arguments"));
647            }
648
649            if spec == 'u' {
650                _warnings::warn(
651                    vm.ctx.exceptions.deprecation_warning,
652                    "The 'u' type code is deprecated and will be removed in Python 3.16".to_owned(),
653                    1,
654                    vm,
655                )?;
656            }
657
658            let mut array =
659                ArrayContentType::from_char(spec).map_err(|err| vm.new_value_error(err))?;
660
661            if let OptionalArg::Present(init) = init {
662                if let Some(init) = init.downcast_ref::<Self>() {
663                    match (spec, init.read().typecode()) {
664                        (spec, ch) if spec == ch => array.frombytes(&init.get_bytes()),
665                        (spec, 'u') => {
666                            return Err(vm.new_type_error(format!(
667                            "cannot use a unicode array to initialize an array with typecode '{spec}'"
668                        )))
669                        }
670                        _ => {
671                            for obj in init.read().iter(vm) {
672                                array.push(obj?, vm)?;
673                            }
674                        }
675                    }
676                } else if let Some(wtf8) = init.downcast_ref::<PyStr>() {
677                    if spec == 'u' {
678                        let bytes = Self::_unicode_to_wchar_bytes(wtf8.as_wtf8(), array.itemsize());
679                        array.frombytes_move(bytes);
680                    } else {
681                        return Err(vm.new_type_error(format!(
682                            "cannot use a str to initialize an array with typecode '{spec}'"
683                        )));
684                    }
685                } else if init.downcastable::<PyBytes>() || init.downcastable::<PyByteArray>() {
686                    init.try_bytes_like(vm, |x| array.frombytes(x))?;
687                } else if let Ok(iter) = ArgIterable::try_from_object(vm, init.clone()) {
688                    for obj in iter.iter(vm)? {
689                        array.push(obj?, vm)?;
690                    }
691                } else {
692                    init.try_bytes_like(vm, |x| array.frombytes(x))?;
693                }
694            }
695
696            Ok(Self::from(array))
697        }
698    }
699
700    #[pyclass(
701        flags(BASETYPE, HAS_WEAKREF),
702        with(
703            Comparable,
704            AsBuffer,
705            AsMapping,
706            AsSequence,
707            Iterable,
708            Constructor,
709            Representable
710        )
711    )]
712    impl PyArray {
713        fn read(&self) -> PyRwLockReadGuard<'_, ArrayContentType> {
714            self.array.read()
715        }
716
717        fn write(&self) -> PyRwLockWriteGuard<'_, ArrayContentType> {
718            self.array.write()
719        }
720
721        #[pygetset]
722        fn typecode(&self, vm: &VirtualMachine) -> PyStrRef {
723            vm.ctx
724                .intern_str(self.read().typecode().to_string())
725                .to_owned()
726        }
727
728        #[pygetset]
729        fn itemsize(&self) -> usize {
730            self.read().itemsize()
731        }
732
733        #[pymethod]
734        fn append(zelf: &Py<Self>, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
735            zelf.try_resizable(vm)?.push(x, vm)
736        }
737
738        #[pymethod]
739        fn clear(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
740            zelf.try_resizable(vm)?.clear()
741        }
742
743        #[pymethod]
744        fn buffer_info(&self) -> (usize, usize) {
745            let array = self.read();
746            (array.addr(), array.len())
747        }
748
749        #[pymethod]
750        fn count(&self, x: PyObjectRef, vm: &VirtualMachine) -> usize {
751            self.read().count(x, vm)
752        }
753
754        #[pymethod]
755        fn remove(zelf: &Py<Self>, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
756            zelf.try_resizable(vm)?.remove(x, vm)
757        }
758
759        #[pymethod]
760        fn extend(zelf: &Py<Self>, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
761            let mut w = zelf.try_resizable(vm)?;
762            if zelf.is(&obj) {
763                w.imul(2, vm)
764            } else if let Some(array) = obj.downcast_ref::<Self>() {
765                w.iadd(&array.read(), vm)
766            } else {
767                let iter = ArgIterable::try_from_object(vm, obj)?;
768                // zelf.extend_from_iterable(iter, vm)
769                for obj in iter.iter(vm)? {
770                    w.push(obj?, vm)?;
771                }
772                Ok(())
773            }
774        }
775
776        fn _wchar_bytes_to_string(
777            bytes: &[u8],
778            item_size: usize,
779            vm: &VirtualMachine,
780        ) -> PyResult<Wtf8Buf> {
781            if item_size == 2 {
782                // safe because every configuration of bytes for the types we support are valid
783                let utf16 = unsafe {
784                    core::slice::from_raw_parts(
785                        bytes.as_ptr() as *const u16,
786                        bytes.len() / core::mem::size_of::<u16>(),
787                    )
788                };
789                Ok(Wtf8Buf::from_wide(utf16))
790            } else {
791                // safe because every configuration of bytes for the types we support are valid
792                let chars = unsafe {
793                    core::slice::from_raw_parts(
794                        bytes.as_ptr() as *const u32,
795                        bytes.len() / core::mem::size_of::<u32>(),
796                    )
797                };
798                chars
799                    .iter()
800                    .map(|&ch| {
801                        // cpython issue 17223
802                        u32_to_char(ch).map_err(|msg| vm.new_value_error(msg))
803                    })
804                    .try_collect()
805            }
806        }
807
808        fn _unicode_to_wchar_bytes(wtf8: &Wtf8, item_size: usize) -> Vec<u8> {
809            if item_size == 2 {
810                wtf8.encode_wide().flat_map(|ch| ch.to_ne_bytes()).collect()
811            } else {
812                wtf8.code_points()
813                    .flat_map(|ch| ch.to_u32().to_ne_bytes())
814                    .collect()
815            }
816        }
817
818        #[pymethod]
819        fn fromunicode(zelf: &Py<Self>, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
820            let wtf8: &Wtf8 = obj.try_to_value(vm).map_err(|_| {
821                vm.new_type_error(format!(
822                    "fromunicode() argument must be str, not {}",
823                    obj.class().name()
824                ))
825            })?;
826            if zelf.read().typecode() != 'u' {
827                return Err(
828                    vm.new_value_error("fromunicode() may only be called on unicode type arrays")
829                );
830            }
831            let mut w = zelf.try_resizable(vm)?;
832            let bytes = Self::_unicode_to_wchar_bytes(wtf8, w.itemsize());
833            w.frombytes_move(bytes);
834            Ok(())
835        }
836
837        #[pymethod]
838        fn tounicode(&self, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
839            let array = self.array.read();
840            if array.typecode() != 'u' {
841                return Err(
842                    vm.new_value_error("tounicode() may only be called on unicode type arrays")
843                );
844            }
845            let bytes = array.get_bytes();
846            Self::_wchar_bytes_to_string(bytes, self.itemsize(), vm)
847        }
848
849        fn _from_bytes(&self, b: &[u8], itemsize: usize, vm: &VirtualMachine) -> PyResult<()> {
850            if !b.len().is_multiple_of(itemsize) {
851                return Err(vm.new_value_error("bytes length not a multiple of item size"));
852            }
853            if b.len() / itemsize > 0 {
854                self.try_resizable(vm)?.frombytes(b);
855            }
856            Ok(())
857        }
858
859        #[pymethod]
860        fn frombytes(&self, b: ArgBytesLike, vm: &VirtualMachine) -> PyResult<()> {
861            let b = b.borrow_buf();
862            let itemsize = self.read().itemsize();
863            self._from_bytes(&b, itemsize, vm)
864        }
865
866        #[pymethod]
867        fn fromfile(&self, f: PyObjectRef, n: isize, vm: &VirtualMachine) -> PyResult<()> {
868            let itemsize = self.itemsize();
869            if n < 0 {
870                return Err(vm.new_value_error("negative count"));
871            }
872            let n = vm.check_repeat_or_overflow_error(itemsize, n)?;
873            let n_bytes = n * itemsize;
874
875            let b = vm.call_method(&f, "read", (n_bytes,))?;
876            let b = b
877                .downcast::<PyBytes>()
878                .map_err(|_| vm.new_type_error("read() didn't return bytes"))?;
879
880            let not_enough_bytes = b.len() != n_bytes;
881
882            self._from_bytes(b.as_bytes(), itemsize, vm)?;
883
884            if not_enough_bytes {
885                Err(vm.new_exception_msg(
886                    vm.ctx.exceptions.eof_error.to_owned(),
887                    "read() didn't return enough bytes".into(),
888                ))
889            } else {
890                Ok(())
891            }
892        }
893
894        #[pymethod]
895        fn byteswap(&self) {
896            self.write().byteswap();
897        }
898
899        #[pymethod]
900        fn index(
901            &self,
902            x: PyObjectRef,
903            range: OptionalRangeArgs,
904            vm: &VirtualMachine,
905        ) -> PyResult<usize> {
906            let (start, stop) = range.saturate(self.__len__(), vm)?;
907            self.read().index(x, start, stop, vm)
908        }
909
910        #[pymethod]
911        fn insert(zelf: &Py<Self>, i: isize, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
912            let mut w = zelf.try_resizable(vm)?;
913            w.insert(i, x, vm)
914        }
915
916        #[pymethod]
917        fn pop(zelf: &Py<Self>, i: OptionalArg<isize>, vm: &VirtualMachine) -> PyResult {
918            let mut w = zelf.try_resizable(vm)?;
919            if w.len() == 0 {
920                Err(vm.new_index_error("pop from empty array"))
921            } else {
922                w.pop(i.unwrap_or(-1), vm)
923            }
924        }
925
926        #[pymethod]
927        pub(crate) fn tobytes(&self) -> Vec<u8> {
928            self.read().get_bytes().to_vec()
929        }
930
931        #[pymethod]
932        fn tofile(&self, f: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
933            /* Write 64K blocks at a time */
934            /* XXX Make the block size settable */
935            const BLOCKSIZE: usize = 64 * 1024;
936
937            let bytes = {
938                let bytes = self.read();
939                bytes.get_bytes().to_vec()
940            };
941
942            for b in bytes.chunks(BLOCKSIZE) {
943                let b = PyBytes::from(b.to_vec()).into_ref(&vm.ctx);
944                vm.call_method(&f, "write", (b,))?;
945            }
946            Ok(())
947        }
948
949        pub(crate) fn get_bytes(&self) -> PyMappedRwLockReadGuard<'_, [u8]> {
950            PyRwLockReadGuard::map(self.read(), |a| a.get_bytes())
951        }
952
953        pub(crate) fn get_bytes_mut(&self) -> PyMappedRwLockWriteGuard<'_, [u8]> {
954            PyRwLockWriteGuard::map(self.write(), |a| a.get_bytes_mut())
955        }
956
957        #[pymethod]
958        fn tolist(&self, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
959            let array = self.read();
960            let mut v = Vec::with_capacity(array.len());
961            for obj in array.iter(vm) {
962                v.push(obj?);
963            }
964            Ok(v)
965        }
966
967        #[pymethod]
968        fn fromlist(zelf: &Py<Self>, list: PyListRef, vm: &VirtualMachine) -> PyResult<()> {
969            zelf.try_resizable(vm)?.fromlist(&list, vm)
970        }
971
972        #[pymethod]
973        fn reverse(&self) {
974            self.write().reverse()
975        }
976
977        #[pymethod]
978        fn __copy__(&self) -> Self {
979            self.array.read().clone().into()
980        }
981
982        #[pymethod]
983        fn __deepcopy__(&self, _memo: PyObjectRef) -> Self {
984            self.__copy__()
985        }
986
987        fn getitem_inner(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult {
988            match SequenceIndex::try_from_borrowed_object(vm, needle, "array")? {
989                SequenceIndex::Int(i) => self.read().getitem_by_index(i, vm),
990                SequenceIndex::Slice(slice) => self.read().getitem_by_slice(slice, vm),
991            }
992        }
993
994        fn __getitem__(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
995            self.getitem_inner(&needle, vm)
996        }
997
998        fn setitem_inner(
999            zelf: &Py<Self>,
1000            needle: &PyObject,
1001            value: PyObjectRef,
1002            vm: &VirtualMachine,
1003        ) -> PyResult<()> {
1004            match SequenceIndex::try_from_borrowed_object(vm, needle, "array")? {
1005                SequenceIndex::Int(i) => zelf.write().setitem_by_index(i, value, vm),
1006                SequenceIndex::Slice(slice) => {
1007                    let cloned;
1008                    let guard;
1009                    let items = if zelf.is(&value) {
1010                        cloned = zelf.read().clone();
1011                        &cloned
1012                    } else {
1013                        match value.downcast_ref::<Self>() {
1014                            Some(array) => {
1015                                guard = array.read();
1016                                &*guard
1017                            }
1018                            None => {
1019                                return Err(vm.new_type_error(format!(
1020                                    "can only assign array (not \"{}\") to array slice",
1021                                    value.class()
1022                                )));
1023                            }
1024                        }
1025                    };
1026                    if let Ok(mut w) = zelf.try_resizable(vm) {
1027                        w.setitem_by_slice(slice, items, vm)
1028                    } else {
1029                        zelf.write().setitem_by_slice_no_resize(slice, items, vm)
1030                    }
1031                }
1032            }
1033        }
1034
1035        fn __setitem__(
1036            zelf: &Py<Self>,
1037            needle: PyObjectRef,
1038            value: PyObjectRef,
1039            vm: &VirtualMachine,
1040        ) -> PyResult<()> {
1041            Self::setitem_inner(zelf, &needle, value, vm)
1042        }
1043
1044        fn delitem_inner(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
1045            match SequenceIndex::try_from_borrowed_object(vm, needle, "array")? {
1046                SequenceIndex::Int(i) => self.try_resizable(vm)?.delitem_by_index(i, vm),
1047                SequenceIndex::Slice(slice) => self.try_resizable(vm)?.delitem_by_slice(slice, vm),
1048            }
1049        }
1050
1051        fn __delitem__(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
1052            self.delitem_inner(&needle, vm)
1053        }
1054
1055        fn __add__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
1056            if let Some(other) = other.downcast_ref::<Self>() {
1057                self.read()
1058                    .add(&other.read(), vm)
1059                    .map(|array| Self::from(array).into_ref(&vm.ctx))
1060            } else {
1061                Err(vm.new_type_error(format!(
1062                    "can only append array (not \"{}\") to array",
1063                    other.class().name()
1064                )))
1065            }
1066        }
1067
1068        fn __iadd__(
1069            zelf: PyRef<Self>,
1070            other: PyObjectRef,
1071            vm: &VirtualMachine,
1072        ) -> PyResult<PyRef<Self>> {
1073            if zelf.is(&other) {
1074                zelf.try_resizable(vm)?.imul(2, vm)?;
1075            } else if let Some(other) = other.downcast_ref::<Self>() {
1076                zelf.try_resizable(vm)?.iadd(&other.read(), vm)?;
1077            } else {
1078                return Err(vm.new_type_error(format!(
1079                    "can only extend array with array (not \"{}\")",
1080                    other.class().name()
1081                )));
1082            }
1083            Ok(zelf)
1084        }
1085
1086        fn __mul__(&self, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
1087            self.read()
1088                .mul(value, vm)
1089                .map(|x| Self::from(x).into_ref(&vm.ctx))
1090        }
1091
1092        fn __imul__(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
1093            zelf.try_resizable(vm)?.imul(value, vm)?;
1094            Ok(zelf)
1095        }
1096
1097        pub(crate) fn __len__(&self) -> usize {
1098            self.read().len()
1099        }
1100
1101        fn array_eq(&self, other: &Self, vm: &VirtualMachine) -> PyResult<bool> {
1102            // we cannot use zelf.is(other) for shortcut because if we contenting a
1103            // float value NaN we always return False even they are the same object.
1104            if self.__len__() != other.__len__() {
1105                return Ok(false);
1106            }
1107            let array_a = self.read();
1108            let array_b = other.read();
1109
1110            // fast path for same ArrayContentType type
1111            if let Ok(ord) = array_a.cmp(&array_b) {
1112                return Ok(ord == Some(Ordering::Equal));
1113            }
1114
1115            let iter = Iterator::zip(array_a.iter(vm), array_b.iter(vm));
1116
1117            for (a, b) in iter {
1118                if !vm.bool_eq(&*a?, &*b?)? {
1119                    return Ok(false);
1120                }
1121            }
1122            Ok(true)
1123        }
1124
1125        #[pymethod]
1126        fn __reduce_ex__(
1127            zelf: &Py<Self>,
1128            proto: usize,
1129            vm: &VirtualMachine,
1130        ) -> PyResult<(PyObjectRef, PyTupleRef, Option<PyDictRef>)> {
1131            if proto < 3 {
1132                return Self::__reduce__(zelf, vm);
1133            }
1134            let array = zelf.read();
1135            let cls = zelf.class().to_owned();
1136            let typecode = vm.ctx.new_str(array.typecode_str());
1137            let bytes = vm.ctx.new_bytes(array.get_bytes().to_vec());
1138            let code = MachineFormatCode::from_typecode(array.typecode()).unwrap();
1139            let code = PyInt::from(u8::from(code)).into_pyobject(vm);
1140            let module = vm.import("array", 0)?;
1141            let func = module.get_attr("_array_reconstructor", vm)?;
1142            Ok((
1143                func,
1144                vm.new_tuple((cls, typecode, code, bytes)),
1145                zelf.as_object().dict(),
1146            ))
1147        }
1148
1149        #[pymethod]
1150        fn __reduce__(
1151            zelf: &Py<Self>,
1152            vm: &VirtualMachine,
1153        ) -> PyResult<(PyObjectRef, PyTupleRef, Option<PyDictRef>)> {
1154            let array = zelf.read();
1155            let cls = zelf.class().to_owned();
1156            let typecode = vm.ctx.new_str(array.typecode_str());
1157            let values = if array.typecode() == 'u' {
1158                let s = Self::_wchar_bytes_to_string(array.get_bytes(), array.itemsize(), vm)?;
1159                s.code_points().map(|x| x.to_pyobject(vm)).collect()
1160            } else {
1161                array.get_objects(vm)
1162            };
1163            let values = vm.ctx.new_list(values);
1164            Ok((
1165                cls.into(),
1166                vm.new_tuple((typecode, values)),
1167                zelf.as_object().dict(),
1168            ))
1169        }
1170
1171        fn __contains__(&self, value: PyObjectRef, vm: &VirtualMachine) -> bool {
1172            let array = self.array.read();
1173            for element in array
1174                .iter(vm)
1175                .map(|x| x.expect("Expected to be checked by array.len() and read lock."))
1176            {
1177                if let Ok(true) =
1178                    element.rich_compare_bool(value.as_object(), PyComparisonOp::Eq, vm)
1179                {
1180                    return true;
1181                }
1182            }
1183
1184            false
1185        }
1186
1187        #[pyclassmethod]
1188        fn __class_getitem__(
1189            cls: PyTypeRef,
1190            args: PyObjectRef,
1191            vm: &VirtualMachine,
1192        ) -> PyGenericAlias {
1193            PyGenericAlias::from_args(cls, args, vm)
1194        }
1195    }
1196
1197    impl Comparable for PyArray {
1198        fn cmp(
1199            zelf: &Py<Self>,
1200            other: &PyObject,
1201            op: PyComparisonOp,
1202            vm: &VirtualMachine,
1203        ) -> PyResult<PyComparisonValue> {
1204            // TODO: deduplicate this logic with sequence::cmp in sequence.rs. Maybe make it generic?
1205
1206            // we cannot use zelf.is(other) for shortcut because if we contenting a
1207            // float value NaN we always return False even they are the same object.
1208            let other = class_or_notimplemented!(Self, other);
1209
1210            if let PyComparisonValue::Implemented(x) =
1211                op.eq_only(|| Ok(zelf.array_eq(other, vm)?.into()))?
1212            {
1213                return Ok(x.into());
1214            }
1215
1216            let array_a = zelf.read();
1217            let array_b = other.read();
1218
1219            let res = match array_a.cmp(&array_b) {
1220                // fast path for same ArrayContentType type
1221                Ok(partial_ord) => partial_ord.is_some_and(|ord| op.eval_ord(ord)),
1222                Err(()) => {
1223                    let iter = Iterator::zip(array_a.iter(vm), array_b.iter(vm));
1224
1225                    for (a, b) in iter {
1226                        let ret = match op {
1227                            PyComparisonOp::Lt | PyComparisonOp::Le => {
1228                                vm.bool_seq_lt(&*a?, &*b?)?
1229                            }
1230                            PyComparisonOp::Gt | PyComparisonOp::Ge => {
1231                                vm.bool_seq_gt(&*a?, &*b?)?
1232                            }
1233                            _ => unreachable!(),
1234                        };
1235                        if let Some(v) = ret {
1236                            return Ok(PyComparisonValue::Implemented(v));
1237                        }
1238                    }
1239
1240                    // fallback:
1241                    op.eval_ord(array_a.len().cmp(&array_b.len()))
1242                }
1243            };
1244
1245            Ok(res.into())
1246        }
1247    }
1248
1249    impl AsBuffer for PyArray {
1250        fn as_buffer(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
1251            let array = zelf.read();
1252            let buf = PyBuffer::new(
1253                zelf.to_owned().into(),
1254                BufferDescriptor::format(
1255                    array.len() * array.itemsize(),
1256                    false,
1257                    array.itemsize(),
1258                    array.typecode_str().into(),
1259                ),
1260                &BUFFER_METHODS,
1261            );
1262            Ok(buf)
1263        }
1264    }
1265
1266    impl Representable for PyArray {
1267        #[inline]
1268        fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
1269            let class = zelf.class();
1270            let class_name = class.name();
1271            if zelf.read().typecode() == 'u' {
1272                if zelf.__len__() == 0 {
1273                    return Ok(format!("{class_name}('u')"));
1274                }
1275                let to_unicode = zelf.tounicode(vm)?;
1276                let escape = crate::vm::literal::escape::UnicodeEscape::new_repr(&to_unicode);
1277                return Ok(format!("{}('u', {})", class_name, escape.str_repr()));
1278            }
1279            zelf.read().repr(&class_name, vm)
1280        }
1281    }
1282
1283    static BUFFER_METHODS: BufferMethods = BufferMethods {
1284        obj_bytes: |buffer| buffer.obj_as::<PyArray>().get_bytes().into(),
1285        obj_bytes_mut: |buffer| buffer.obj_as::<PyArray>().get_bytes_mut().into(),
1286        release: |buffer| {
1287            buffer
1288                .obj_as::<PyArray>()
1289                .exports
1290                .fetch_sub(1, atomic::Ordering::Release);
1291        },
1292        retain: |buffer| {
1293            buffer
1294                .obj_as::<PyArray>()
1295                .exports
1296                .fetch_add(1, atomic::Ordering::Release);
1297        },
1298    };
1299
1300    impl AsMapping for PyArray {
1301        fn as_mapping() -> &'static PyMappingMethods {
1302            static AS_MAPPING: PyMappingMethods = PyMappingMethods {
1303                length: atomic_func!(|mapping, _vm| Ok(
1304                    PyArray::mapping_downcast(mapping).__len__()
1305                )),
1306                subscript: atomic_func!(|mapping, needle, vm| {
1307                    PyArray::mapping_downcast(mapping).getitem_inner(needle, vm)
1308                }),
1309                ass_subscript: atomic_func!(|mapping, needle, value, vm| {
1310                    let zelf = PyArray::mapping_downcast(mapping);
1311                    if let Some(value) = value {
1312                        PyArray::setitem_inner(zelf, needle, value, vm)
1313                    } else {
1314                        zelf.delitem_inner(needle, vm)
1315                    }
1316                }),
1317            };
1318            &AS_MAPPING
1319        }
1320    }
1321
1322    impl AsSequence for PyArray {
1323        fn as_sequence() -> &'static PySequenceMethods {
1324            static AS_SEQUENCE: PySequenceMethods = PySequenceMethods {
1325                length: atomic_func!(|seq, _vm| Ok(PyArray::sequence_downcast(seq).__len__())),
1326                concat: atomic_func!(|seq, other, vm| {
1327                    let zelf = PyArray::sequence_downcast(seq);
1328                    PyArray::__add__(zelf, other.to_owned(), vm).map(|x| x.into())
1329                }),
1330                repeat: atomic_func!(|seq, n, vm| {
1331                    PyArray::sequence_downcast(seq)
1332                        .__mul__(n, vm)
1333                        .map(|x| x.into())
1334                }),
1335                item: atomic_func!(|seq, i, vm| {
1336                    PyArray::sequence_downcast(seq)
1337                        .read()
1338                        .getitem_by_index(i, vm)
1339                }),
1340                ass_item: atomic_func!(|seq, i, value, vm| {
1341                    let zelf = PyArray::sequence_downcast(seq);
1342                    if let Some(value) = value {
1343                        zelf.write().setitem_by_index(i, value, vm)
1344                    } else {
1345                        zelf.write().delitem_by_index(i, vm)
1346                    }
1347                }),
1348                contains: atomic_func!(|seq, target, vm| {
1349                    let zelf = PyArray::sequence_downcast(seq);
1350                    Ok(zelf.__contains__(target.to_owned(), vm))
1351                }),
1352                inplace_concat: atomic_func!(|seq, other, vm| {
1353                    let zelf = PyArray::sequence_downcast(seq).to_owned();
1354                    PyArray::__iadd__(zelf, other.to_owned(), vm).map(|x| x.into())
1355                }),
1356                inplace_repeat: atomic_func!(|seq, n, vm| {
1357                    let zelf = PyArray::sequence_downcast(seq).to_owned();
1358                    PyArray::__imul__(zelf, n, vm).map(|x| x.into())
1359                }),
1360            };
1361            &AS_SEQUENCE
1362        }
1363    }
1364
1365    impl Iterable for PyArray {
1366        fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
1367            Ok(PyArrayIter {
1368                internal: PyMutex::new(PositionIterInternal::new(zelf, 0)),
1369            }
1370            .into_pyobject(vm))
1371        }
1372    }
1373
1374    impl BufferResizeGuard for PyArray {
1375        type Resizable<'a> = PyRwLockWriteGuard<'a, ArrayContentType>;
1376
1377        fn try_resizable_opt(&self) -> Option<Self::Resizable<'_>> {
1378            let w = self.write();
1379            (self.exports.load(atomic::Ordering::SeqCst) == 0).then_some(w)
1380        }
1381    }
1382
1383    #[pyattr]
1384    #[pyclass(name = "arrayiterator", traverse)]
1385    #[derive(Debug, PyPayload)]
1386    pub struct PyArrayIter {
1387        internal: PyMutex<PositionIterInternal<PyArrayRef>>,
1388    }
1389
1390    #[pyclass(with(IterNext, Iterable), flags(HAS_DICT, DISALLOW_INSTANTIATION))]
1391    impl PyArrayIter {
1392        #[pymethod]
1393        fn __setstate__(&self, state: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
1394            self.internal
1395                .lock()
1396                .set_state(state, |obj, pos| pos.min(obj.__len__()), vm)
1397        }
1398
1399        #[pymethod]
1400        fn __reduce__(&self, vm: &VirtualMachine) -> PyTupleRef {
1401            let func = builtins_iter(vm);
1402            self.internal.lock().reduce(
1403                func,
1404                |x| x.clone().into(),
1405                |vm| vm.ctx.empty_tuple.clone().into(),
1406                vm,
1407            )
1408        }
1409    }
1410
1411    impl SelfIter for PyArrayIter {}
1412    impl IterNext for PyArrayIter {
1413        fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
1414            zelf.internal.lock().next(|array, pos| {
1415                Ok(match array.read().get(pos, vm) {
1416                    Some(item) => PyIterReturn::Return(item?),
1417                    None => PyIterReturn::StopIteration(None),
1418                })
1419            })
1420        }
1421    }
1422
1423    #[derive(FromArgs)]
1424    struct ReconstructorArgs {
1425        #[pyarg(positional)]
1426        arraytype: PyTypeRef,
1427        #[pyarg(positional)]
1428        typecode: PyUtf8StrRef,
1429        #[pyarg(positional)]
1430        mformat_code: MachineFormatCode,
1431        #[pyarg(positional)]
1432        items: PyBytesRef,
1433    }
1434
1435    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
1436    #[repr(u8)]
1437    enum MachineFormatCode {
1438        Int8 { signed: bool },                    // 0, 1
1439        Int16 { signed: bool, big_endian: bool }, // 2, 3, 4, 5
1440        Int32 { signed: bool, big_endian: bool }, // 6, 7, 8, 9
1441        Int64 { signed: bool, big_endian: bool }, // 10, 11, 12, 13
1442        Ieee754Float { big_endian: bool },        // 14, 15
1443        Ieee754Double { big_endian: bool },       // 16, 17
1444        Utf16 { big_endian: bool },               // 18, 19
1445        Utf32 { big_endian: bool },               // 20, 21
1446    }
1447
1448    impl From<MachineFormatCode> for u8 {
1449        fn from(code: MachineFormatCode) -> Self {
1450            use MachineFormatCode::*;
1451            match code {
1452                Int8 { signed } => signed as Self,
1453                Int16 { signed, big_endian } => 2 + signed as Self * 2 + big_endian as Self,
1454                Int32 { signed, big_endian } => 6 + signed as Self * 2 + big_endian as Self,
1455                Int64 { signed, big_endian } => 10 + signed as Self * 2 + big_endian as Self,
1456                Ieee754Float { big_endian } => 14 + big_endian as Self,
1457                Ieee754Double { big_endian } => 16 + big_endian as Self,
1458                Utf16 { big_endian } => 18 + big_endian as Self,
1459                Utf32 { big_endian } => 20 + big_endian as Self,
1460            }
1461        }
1462    }
1463
1464    impl TryFrom<u8> for MachineFormatCode {
1465        type Error = u8;
1466
1467        fn try_from(code: u8) -> Result<Self, Self::Error> {
1468            let big_endian = !code.is_multiple_of(2);
1469            let signed = match code {
1470                0 | 1 => code != 0,
1471                2..=13 => (code - 2) % 4 >= 2,
1472                _ => false,
1473            };
1474            match code {
1475                0..=1 => Ok(Self::Int8 { signed }),
1476                2..=5 => Ok(Self::Int16 { signed, big_endian }),
1477                6..=9 => Ok(Self::Int32 { signed, big_endian }),
1478                10..=13 => Ok(Self::Int64 { signed, big_endian }),
1479                14..=15 => Ok(Self::Ieee754Float { big_endian }),
1480                16..=17 => Ok(Self::Ieee754Double { big_endian }),
1481                18..=19 => Ok(Self::Utf16 { big_endian }),
1482                20..=21 => Ok(Self::Utf32 { big_endian }),
1483                _ => Err(code),
1484            }
1485        }
1486    }
1487
1488    impl<'a> TryFromBorrowedObject<'a> for MachineFormatCode {
1489        fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {
1490            obj.try_to_ref::<PyInt>(vm)
1491                .map_err(|_| {
1492                    vm.new_type_error(format!(
1493                        "an integer is required (got type {})",
1494                        obj.class().name()
1495                    ))
1496                })?
1497                .try_to_primitive::<i32>(vm)?
1498                .to_u8()
1499                .unwrap_or(u8::MAX)
1500                .try_into()
1501                .map_err(|_| {
1502                    vm.new_value_error("third argument must be a valid machine format code.")
1503                })
1504        }
1505    }
1506
1507    impl MachineFormatCode {
1508        fn from_typecode(code: char) -> Option<Self> {
1509            use core::mem::size_of;
1510            let signed = code.is_ascii_uppercase();
1511            let big_endian = cfg!(target_endian = "big");
1512            let int_size = match code {
1513                'b' | 'B' => return Some(Self::Int8 { signed }),
1514                'u' => {
1515                    return match size_of::<wchar_t>() {
1516                        2 => Some(Self::Utf16 { big_endian }),
1517                        4 => Some(Self::Utf32 { big_endian }),
1518                        _ => None,
1519                    };
1520                }
1521                'f' => {
1522                    // Copied from CPython
1523                    const Y: f32 = 16711938.0;
1524                    return match &Y.to_ne_bytes() {
1525                        b"\x4b\x7f\x01\x02" => Some(Self::Ieee754Float { big_endian: true }),
1526                        b"\x02\x01\x7f\x4b" => Some(Self::Ieee754Float { big_endian: false }),
1527                        _ => None,
1528                    };
1529                }
1530                'd' => {
1531                    // Copied from CPython
1532                    const Y: f64 = 9006104071832581.0;
1533                    return match &Y.to_ne_bytes() {
1534                        b"\x43\x3f\xff\x01\x02\x03\x04\x05" => {
1535                            Some(Self::Ieee754Double { big_endian: true })
1536                        }
1537                        b"\x05\x04\x03\x02\x01\xff\x3f\x43" => {
1538                            Some(Self::Ieee754Double { big_endian: false })
1539                        }
1540                        _ => None,
1541                    };
1542                }
1543                _ => ArrayContentType::itemsize_of_typecode(code)? as u8,
1544            };
1545            match int_size {
1546                2 => Some(Self::Int16 { signed, big_endian }),
1547                4 => Some(Self::Int32 { signed, big_endian }),
1548                8 => Some(Self::Int64 { signed, big_endian }),
1549                _ => None,
1550            }
1551        }
1552        const fn item_size(self) -> usize {
1553            match self {
1554                Self::Int8 { .. } => 1,
1555                Self::Int16 { .. } | Self::Utf16 { .. } => 2,
1556                Self::Int32 { .. } | Self::Utf32 { .. } | Self::Ieee754Float { .. } => 4,
1557                Self::Int64 { .. } | Self::Ieee754Double { .. } => 8,
1558            }
1559        }
1560    }
1561
1562    fn check_array_type(typ: PyTypeRef, vm: &VirtualMachine) -> PyResult<PyTypeRef> {
1563        if !typ.fast_issubclass(PyArray::class(&vm.ctx)) {
1564            return Err(
1565                vm.new_type_error(format!("{} is not a subtype of array.array", typ.name()))
1566            );
1567        }
1568        Ok(typ)
1569    }
1570
1571    fn check_type_code(spec: PyUtf8StrRef, vm: &VirtualMachine) -> PyResult<ArrayContentType> {
1572        let spec = spec.as_str().chars().exactly_one().map_err(|_| {
1573            vm.new_type_error(
1574                "_array_reconstructor() argument 2 must be a unicode character, not str",
1575            )
1576        })?;
1577        ArrayContentType::from_char(spec)
1578            .map_err(|_| vm.new_value_error("second argument must be a valid type code"))
1579    }
1580
1581    macro_rules! chunk_to_obj {
1582        ($BYTE:ident, $TY:ty, $BIG_ENDIAN:ident) => {{
1583            let b = <[u8; ::core::mem::size_of::<$TY>()]>::try_from($BYTE).unwrap();
1584            if $BIG_ENDIAN {
1585                <$TY>::from_be_bytes(b)
1586            } else {
1587                <$TY>::from_le_bytes(b)
1588            }
1589        }};
1590        ($VM:ident, $BYTE:ident, $TY:ty, $BIG_ENDIAN:ident) => {
1591            chunk_to_obj!($BYTE, $TY, $BIG_ENDIAN).to_pyobject($VM)
1592        };
1593        ($VM:ident, $BYTE:ident, $SIGNED_TY:ty, $UNSIGNED_TY:ty, $SIGNED:ident, $BIG_ENDIAN:ident) => {{
1594            let b = <[u8; ::core::mem::size_of::<$SIGNED_TY>()]>::try_from($BYTE).unwrap();
1595            match ($SIGNED, $BIG_ENDIAN) {
1596                (false, false) => <$UNSIGNED_TY>::from_le_bytes(b).to_pyobject($VM),
1597                (false, true) => <$UNSIGNED_TY>::from_be_bytes(b).to_pyobject($VM),
1598                (true, false) => <$SIGNED_TY>::from_le_bytes(b).to_pyobject($VM),
1599                (true, true) => <$SIGNED_TY>::from_be_bytes(b).to_pyobject($VM),
1600            }
1601        }};
1602    }
1603
1604    #[pyfunction]
1605    fn _array_reconstructor(args: ReconstructorArgs, vm: &VirtualMachine) -> PyResult<PyArrayRef> {
1606        let cls = check_array_type(args.arraytype, vm)?;
1607        let mut array = check_type_code(args.typecode, vm)?;
1608        let format = args.mformat_code;
1609        let bytes = args.items.as_bytes();
1610        if !bytes.len().is_multiple_of(format.item_size()) {
1611            return Err(vm.new_value_error("bytes length not a multiple of item size"));
1612        }
1613        if MachineFormatCode::from_typecode(array.typecode()) == Some(format) {
1614            array.frombytes(bytes);
1615            return PyArray::from(array).into_ref_with_type(vm, cls);
1616        }
1617        if !matches!(
1618            format,
1619            MachineFormatCode::Utf16 { .. } | MachineFormatCode::Utf32 { .. }
1620        ) {
1621            array.reserve(bytes.len() / format.item_size());
1622        }
1623        let mut chunks = bytes.chunks(format.item_size());
1624        match format {
1625            MachineFormatCode::Ieee754Float { big_endian } => {
1626                chunks.try_for_each(|b| array.push(chunk_to_obj!(vm, b, f32, big_endian), vm))?
1627            }
1628            MachineFormatCode::Ieee754Double { big_endian } => {
1629                chunks.try_for_each(|b| array.push(chunk_to_obj!(vm, b, f64, big_endian), vm))?
1630            }
1631            MachineFormatCode::Int8 { signed } => chunks
1632                .try_for_each(|b| array.push(chunk_to_obj!(vm, b, i8, u8, signed, false), vm))?,
1633            MachineFormatCode::Int16 { signed, big_endian } => chunks.try_for_each(|b| {
1634                array.push(chunk_to_obj!(vm, b, i16, u16, signed, big_endian), vm)
1635            })?,
1636            MachineFormatCode::Int32 { signed, big_endian } => chunks.try_for_each(|b| {
1637                array.push(chunk_to_obj!(vm, b, i32, u32, signed, big_endian), vm)
1638            })?,
1639            MachineFormatCode::Int64 { signed, big_endian } => chunks.try_for_each(|b| {
1640                array.push(chunk_to_obj!(vm, b, i64, u64, signed, big_endian), vm)
1641            })?,
1642            MachineFormatCode::Utf16 { big_endian } => {
1643                let utf16: Vec<_> = chunks.map(|b| chunk_to_obj!(b, u16, big_endian)).collect();
1644                let s = String::from_utf16(&utf16)
1645                    .map_err(|_| vm.new_unicode_encode_error("items cannot decode as utf16"))?;
1646                let bytes = PyArray::_unicode_to_wchar_bytes((*s).as_ref(), array.itemsize());
1647                array.frombytes_move(bytes);
1648            }
1649            MachineFormatCode::Utf32 { big_endian } => {
1650                let s: Wtf8Buf = chunks
1651                    .map(|b| chunk_to_obj!(b, u32, big_endian))
1652                    .map(|ch| u32_to_char(ch).map_err(|msg| vm.new_value_error(msg)))
1653                    .try_collect()?;
1654                let bytes = PyArray::_unicode_to_wchar_bytes(&s, array.itemsize());
1655                array.frombytes_move(bytes);
1656            }
1657        };
1658        PyArray::from(array).into_ref_with_type(vm, cls)
1659    }
1660
1661    // Register array.array as collections.abc.MutableSequence
1662    pub(crate) fn module_exec(
1663        vm: &VirtualMachine,
1664        module: &Py<crate::vm::builtins::PyModule>,
1665    ) -> PyResult<()> {
1666        __module_exec(vm, module);
1667
1668        let array_type = module
1669            .get_attr("array", vm)
1670            .expect("array module has array type");
1671
1672        // vm.import returns the top-level module, so we need to get abc submodule
1673        let collections_abc = vm.import("collections.abc", 0)?;
1674        let abc = collections_abc.get_attr("abc", vm)?;
1675        let mutable_sequence = abc.get_attr("MutableSequence", vm)?;
1676        let register = mutable_sequence.get_attr("register", vm)?;
1677        register.call((array_type,), vm)?;
1678
1679        Ok(())
1680    }
1681}