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