1pub(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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 if self.__len__() != other.__len__() {
1105 return Ok(false);
1106 }
1107 let array_a = self.read();
1108 let array_b = other.read();
1109
1110 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 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 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 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 }, 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 }, }
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 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 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 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 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}