Merge pull request RustPython#4677 from youknowone/arg-index

apply ArgIndex

Author: youknowone

stdlib/src/array.rs

@@ -509,7 +509,23 @@ mod array {
         fn to_object(self, vm: &VirtualMachine) -> PyObjectRef;
     }
 
-    macro_rules! impl_array_element {
+    macro_rules! impl_int_element {
+        ($($t:ty,)*) => {$(
+            impl ArrayElement for $t {
+                fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+                    obj.try_index(vm)?.try_to_primitive(vm)
+                }
+                fn byteswap(self) -> Self {
+                    <$t>::swap_bytes(self)
+                }
+                fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {
+                    self.to_pyobject(vm)
+                }
+            }
+        )*};
+    }
+
+    macro_rules! impl_float_element {
         ($(($t:ty, $f_from:path, $f_swap:path, $f_to:path),)*) => {$(
             impl ArrayElement for $t {
                 fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
@@ -525,15 +541,8 @@ mod array {
         )*};
     }
 
-    impl_array_element!(
-        (i8, i8::try_from_object, i8::swap_bytes, PyInt::from),
-        (u8, u8::try_from_object, u8::swap_bytes, PyInt::from),
-        (i16, i16::try_from_object, i16::swap_bytes, PyInt::from),
-        (u16, u16::try_from_object, u16::swap_bytes, PyInt::from),
-        (i32, i32::try_from_object, i32::swap_bytes, PyInt::from),
-        (u32, u32::try_from_object, u32::swap_bytes, PyInt::from),
-        (i64, i64::try_from_object, i64::swap_bytes, PyInt::from),
-        (u64, u64::try_from_object, u64::swap_bytes, PyInt::from),
+    impl_int_element!(i8, u8, i16, u16, i32, u32, i64, u64,);
+    impl_float_element!(
         (
             f32,
             f32_try_into_from_object,

stdlib/src/bisect.rs

@@ -3,29 +3,28 @@ pub(crate) use _bisect::make_module;
 #[pymodule]
 mod _bisect {
     use crate::vm::{
-        function::OptionalArg, types::PyComparisonOp, PyObjectRef, PyResult, VirtualMachine,
+        function::{ArgIndex, OptionalArg},
+        types::PyComparisonOp,
+        PyObjectRef, PyResult, VirtualMachine,
     };
 
     #[derive(FromArgs)]
     struct BisectArgs {
         a: PyObjectRef,
         x: PyObjectRef,
         #[pyarg(any, optional)]
-        lo: OptionalArg<PyObjectRef>,
+        lo: OptionalArg<ArgIndex>,
         #[pyarg(any, optional)]
-        hi: OptionalArg<PyObjectRef>,
+        hi: OptionalArg<ArgIndex>,
         #[pyarg(named, default)]
         key: Option<PyObjectRef>,
     }
 
     // Handles objects that implement __index__ and makes sure index fits in needed isize.
     #[inline]
-    fn handle_default(
-        arg: OptionalArg<PyObjectRef>,
-        vm: &VirtualMachine,
-    ) -> PyResult<Option<isize>> {
+    fn handle_default(arg: OptionalArg<ArgIndex>, vm: &VirtualMachine) -> PyResult<Option<isize>> {
         arg.into_option()
-            .map(|v| v.try_index(vm)?.try_to_primitive(vm))
+            .map(|v| v.try_to_primitive(vm))
             .transpose()
     }
 
@@ -38,8 +37,8 @@ mod _bisect {
     //    input sequence.
     #[inline]
     fn as_usize(
-        lo: OptionalArg<PyObjectRef>,
-        hi: OptionalArg<PyObjectRef>,
+        lo: OptionalArg<ArgIndex>,
+        hi: OptionalArg<ArgIndex>,
         seq_len: usize,
         vm: &VirtualMachine,
     ) -> PyResult<(usize, usize)> {

stdlib/src/math.rs

@@ -4,7 +4,7 @@ pub(crate) use math::make_module;
 mod math {
     use crate::vm::{
         builtins::{try_bigint_to_f64, try_f64_to_bigint, PyFloat, PyInt, PyIntRef, PyStrInterned},
-        function::{ArgIntoFloat, ArgIterable, Either, OptionalArg, PosArgs},
+        function::{ArgIndex, ArgIntoFloat, ArgIterable, Either, OptionalArg, PosArgs},
         identifier, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine,
     };
     use itertools::Itertools;
@@ -218,9 +218,8 @@ mod math {
     }
 
     #[pyfunction]
-    fn isqrt(x: PyObjectRef, vm: &VirtualMachine) -> PyResult<BigInt> {
-        let index = x.try_index(vm)?;
-        let value = index.as_bigint();
+    fn isqrt(x: ArgIndex, vm: &VirtualMachine) -> PyResult<BigInt> {
+        let value = x.as_bigint();
 
         if value.is_negative() {
             return Err(vm.new_value_error("isqrt() argument must be nonnegative".to_owned()));
@@ -575,7 +574,7 @@ mod math {
         }
     }
 
-    fn math_perf_arb_len_int_op<F>(args: PosArgs<PyIntRef>, op: F, default: BigInt) -> BigInt
+    fn math_perf_arb_len_int_op<F>(args: PosArgs<ArgIndex>, op: F, default: BigInt) -> BigInt
     where
         F: Fn(&BigInt, &PyInt) -> BigInt,
     {
@@ -595,13 +594,13 @@ mod math {
     }
 
     #[pyfunction]
-    fn gcd(args: PosArgs<PyIntRef>) -> BigInt {
+    fn gcd(args: PosArgs<ArgIndex>) -> BigInt {
         use num_integer::Integer;
         math_perf_arb_len_int_op(args, |x, y| x.gcd(y.as_bigint()), BigInt::zero())
     }
 
     #[pyfunction]
-    fn lcm(args: PosArgs<PyIntRef>) -> BigInt {
+    fn lcm(args: PosArgs<ArgIndex>) -> BigInt {
         use num_integer::Integer;
         math_perf_arb_len_int_op(args, |x, y| x.lcm(y.as_bigint()), BigInt::one())
     }
@@ -733,8 +732,8 @@ mod math {
 
     #[pyfunction]
     fn perm(
-        n: PyIntRef,
-        k: OptionalArg<Option<PyIntRef>>,
+        n: ArgIndex,
+        k: OptionalArg<Option<ArgIndex>>,
         vm: &VirtualMachine,
     ) -> PyResult<BigInt> {
         let n = n.as_bigint();
@@ -764,7 +763,7 @@ mod math {
     }
 
     #[pyfunction]
-    fn comb(n: PyIntRef, k: PyIntRef, vm: &VirtualMachine) -> PyResult<BigInt> {
+    fn comb(n: ArgIndex, k: ArgIndex, vm: &VirtualMachine) -> PyResult<BigInt> {
         let mut k = k.as_bigint();
         let n = n.as_bigint();
         let one = BigInt::one();

stdlib/src/zlib.rs

@@ -2,10 +2,10 @@ pub(crate) use zlib::make_module;
 
 #[pymodule]
 mod zlib {
-    use crate::common::lock::PyMutex;
     use crate::vm::{
         builtins::{PyBaseExceptionRef, PyBytes, PyBytesRef, PyIntRef, PyTypeRef},
-        function::{ArgBytesLike, OptionalArg, OptionalOption},
+        common::lock::PyMutex,
+        function::{ArgBytesLike, ArgPrimitiveIndex, ArgSize, OptionalArg, OptionalOption},
         PyPayload, PyResult, VirtualMachine,
     };
     use adler32::RollingAdler32 as Adler32;
@@ -233,9 +233,9 @@ mod zlib {
         #[pyarg(positional)]
         data: ArgBytesLike,
         #[pyarg(any, optional)]
-        wbits: OptionalArg<i8>,
+        wbits: OptionalArg<ArgPrimitiveIndex<i8>>,
         #[pyarg(any, optional)]
-        bufsize: OptionalArg<usize>,
+        bufsize: OptionalArg<ArgPrimitiveIndex<usize>>,
     }
 
     /// Returns a bytes object containing the uncompressed data.
@@ -245,9 +245,9 @@ mod zlib {
         let wbits = arg.wbits;
         let bufsize = arg.bufsize;
         data.with_ref(|data| {
-            let bufsize = bufsize.unwrap_or(DEF_BUF_SIZE);
+            let bufsize = bufsize.into_primitive().unwrap_or(DEF_BUF_SIZE);
 
-            let mut d = header_from_wbits(wbits, vm)?.decompress();
+            let mut d = header_from_wbits(wbits.into_primitive(), vm)?.decompress();
 
             _decompress(data, &mut d, bufsize, None, false, vm).and_then(|(buf, stream_end)| {
                 if stream_end {
@@ -265,7 +265,7 @@ mod zlib {
     #[pyfunction]
     fn decompressobj(args: DecompressobjArgs, vm: &VirtualMachine) -> PyResult<PyDecompress> {
         #[allow(unused_mut)]
-        let mut decompress = header_from_wbits(args.wbits, vm)?.decompress();
+        let mut decompress = header_from_wbits(args.wbits.into_primitive(), vm)?.decompress();
         #[cfg(feature = "zlib")]
         if let OptionalArg::Present(dict) = args.zdict {
             dict.with_ref(|d| decompress.set_dictionary(d).unwrap());
@@ -325,11 +325,8 @@ mod zlib {
 
         #[pymethod]
         fn decompress(&self, args: DecompressArgs, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
-            let max_length = if args.max_length == 0 {
-                None
-            } else {
-                Some(args.max_length)
-            };
+            let max_length = args.max_length.value;
+            let max_length = (max_length != 0).then_some(max_length);
             let data = args.data.borrow_buf();
             let data = &*data;
 
@@ -362,12 +359,18 @@ mod zlib {
         }
 
         #[pymethod]
-        fn flush(&self, length: OptionalArg<isize>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+        fn flush(&self, length: OptionalArg<ArgSize>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
             let length = match length {
-                OptionalArg::Present(l) if l <= 0 => {
-                    return Err(vm.new_value_error("length must be greater than zero".to_owned()));
+                OptionalArg::Present(l) => {
+                    let l: isize = l.into();
+                    if l <= 0 {
+                        return Err(
+                            vm.new_value_error("length must be greater than zero".to_owned())
+                        );
+                    } else {
+                        l as usize
+                    }
                 }
-                OptionalArg::Present(l) => l as usize,
                 OptionalArg::Missing => DEF_BUF_SIZE,
             };
 
@@ -396,14 +399,17 @@ mod zlib {
     struct DecompressArgs {
         #[pyarg(positional)]
         data: ArgBytesLike,
-        #[pyarg(any, default = "0")]
-        max_length: usize,
+        #[pyarg(
+            any,
+            default = "rustpython_vm::function::ArgPrimitiveIndex { value: 0 }"
+        )]
+        max_length: ArgPrimitiveIndex<usize>,
     }
 
     #[derive(FromArgs)]
     struct DecompressobjArgs {
         #[pyarg(any, optional)]
-        wbits: OptionalArg<i8>,
+        wbits: OptionalArg<ArgPrimitiveIndex<i8>>,
         #[cfg(feature = "zlib")]
         #[pyarg(any, optional)]
         zdict: OptionalArg<ArgBytesLike>,
@@ -414,7 +420,7 @@ mod zlib {
         level: OptionalArg<i32>,
         // only DEFLATED is valid right now, it's w/e
         _method: OptionalArg<i32>,
-        wbits: OptionalArg<i8>,
+        wbits: OptionalArg<ArgPrimitiveIndex<i8>>,
         // these aren't used.
         _mem_level: OptionalArg<i32>, // this is memLevel in CPython
         _strategy: OptionalArg<i32>,
@@ -423,7 +429,7 @@ mod zlib {
     ) -> PyResult<PyCompress> {
         let level = compression_from_int(level.into_option())
             .ok_or_else(|| vm.new_value_error("invalid initialization option".to_owned()))?;
-        let compress = header_from_wbits(wbits, vm)?.compress(level);
+        let compress = header_from_wbits(wbits.into_primitive(), vm)?.compress(level);
         Ok(PyCompress {
             inner: PyMutex::new(CompressInner {
                 compress,

vm/src/builtins/bytearray.rs

@@ -20,9 +20,9 @@ use crate::{
         },
     },
     convert::{ToPyObject, ToPyResult},
-    function::Either,
     function::{
-        ArgBytesLike, ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue,
+        ArgBytesLike, ArgIterable, ArgSize, Either, FuncArgs, OptionalArg, OptionalOption,
+        PyComparisonValue,
     },
     protocol::{
         BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
@@ -67,6 +67,10 @@ impl PyByteArray {
     pub fn borrow_buf_mut(&self) -> PyMappedRwLockWriteGuard<'_, Vec<u8>> {
         PyRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements)
     }
+
+    fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
+        self.inner().mul(value, vm).map(|x| x.into())
+    }
 }
 
 impl From<PyBytesInner> for PyByteArray {
@@ -658,13 +662,13 @@ impl PyByteArray {
 
     #[pymethod(name = "__rmul__")]
     #[pymethod(magic)]
-    fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
-        self.inner().mul(value, vm).map(|x| x.into())
+    fn mul(&self, value: ArgSize, vm: &VirtualMachine) -> PyResult<Self> {
+        self.repeat(value.into(), vm)
     }
 
     #[pymethod(magic)]
-    fn imul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-        Self::irepeat(&zelf, value, vm)?;
+    fn imul(zelf: PyRef<Self>, value: ArgSize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        Self::irepeat(&zelf, value.into(), vm)?;
         Ok(zelf)
     }
 
@@ -820,7 +824,7 @@ impl AsSequence for PyByteArray {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 PyByteArray::sequence_downcast(seq)
-                    .mul(n, vm)
+                    .repeat(n, vm)
                     .map(|x| x.into_pyobject(vm))
             }),
             item: atomic_func!(|seq, i, vm| {
@@ -849,7 +853,8 @@ impl AsSequence for PyByteArray {
             }),
             inplace_repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyByteArray::sequence_downcast(seq).to_owned();
-                PyByteArray::imul(zelf, n, vm).map(|x| x.into())
+                PyByteArray::irepeat(&zelf, n, vm)?;
+                Ok(zelf.into())
             }),
         };
         &AS_SEQUENCE

vm/src/builtins/bytes.rs

@@ -11,8 +11,9 @@ use crate::{
     class::PyClassImpl,
     common::{hash::PyHash, lock::PyMutex},
     convert::{ToPyObject, ToPyResult},
-    function::Either,
-    function::{ArgBytesLike, ArgIterable, OptionalArg, OptionalOption, PyComparisonValue},
+    function::{
+        ArgBytesLike, ArgIndex, ArgIterable, Either, OptionalArg, OptionalOption, PyComparisonValue,
+    },
     protocol::{
         BufferDescriptor, BufferMethods, PyBuffer, PyIterReturn, PyMappingMethods, PyNumberMethods,
         PySequenceMethods,
@@ -99,6 +100,19 @@ impl PyBytes {
     pub fn new_ref(data: Vec<u8>, ctx: &Context) -> PyRef<Self> {
         PyRef::new_ref(Self::from(data), ctx.types.bytes_type.to_owned(), None)
     }
+
+    fn repeat(zelf: PyRef<Self>, count: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        if count == 1 && zelf.class().is(vm.ctx.types.bytes_type) {
+            // Special case: when some `bytes` is multiplied by `1`,
+            // nothing really happens, we need to return an object itself
+            // with the same `id()` to be compatible with CPython.
+            // This only works for `bytes` itself, not its subclasses.
+            return Ok(zelf);
+        }
+        zelf.inner
+            .mul(count, vm)
+            .map(|x| Self::from(x).into_ref(vm))
+    }
 }
 
 #[pyclass(
@@ -497,17 +511,8 @@ impl PyBytes {
 
     #[pymethod(name = "__rmul__")]
     #[pymethod(magic)]
-    fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-        if value == 1 && zelf.class().is(vm.ctx.types.bytes_type) {
-            // Special case: when some `bytes` is multiplied by `1`,
-            // nothing really happens, we need to return an object itself
-            // with the same `id()` to be compatible with CPython.
-            // This only works for `bytes` itself, not its subclasses.
-            return Ok(zelf);
-        }
-        zelf.inner
-            .mul(value, vm)
-            .map(|x| Self::from(x).into_ref(vm))
+    fn mul(zelf: PyRef<Self>, value: ArgIndex, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        Self::repeat(zelf, value.try_to_primitive(vm)?, vm)
     }
 
     #[pymethod(name = "__mod__")]
@@ -605,7 +610,7 @@ impl AsSequence for PyBytes {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 if let Ok(zelf) = seq.obj.to_owned().downcast::<PyBytes>() {
-                    PyBytes::mul(zelf, n, vm).to_pyresult(vm)
+                    PyBytes::repeat(zelf, n, vm).to_pyresult(vm)
                 } else {
                     Err(vm.new_type_error("bad argument type for built-in operation".to_owned()))
                 }