ArgPrimitiveIndex for zlib

Author: youknowone

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/str.rs

@@ -13,7 +13,7 @@ use crate::{
     },
     convert::{IntoPyException, ToPyException, ToPyObject, ToPyResult},
     format::{format, format_map},
-    function::{ArgSize, ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
+    function::{ArgIterable, ArgSize, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
     intern::PyInterned,
     protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
     sequence::SequenceExt,

vm/src/function/mod.rs

@@ -17,7 +17,7 @@ pub use builtin::{IntoPyNativeFunc, OwnedParam, PyNativeFunc, RefParam};
 pub use either::Either;
 pub use getset::PySetterValue;
 pub(super) use getset::{IntoPyGetterFunc, IntoPySetterFunc, PyGetterFunc, PySetterFunc};
-pub use number::{ArgIndex, ArgIntoBool, ArgIntoComplex, ArgIntoFloat, ArgSize};
+pub use number::{ArgIndex, ArgIntoBool, ArgIntoComplex, ArgIntoFloat, ArgPrimitiveIndex, ArgSize};
 pub use protocol::{ArgCallable, ArgIterable, ArgMapping, ArgSequence};
 
 use crate::{builtins::PyStr, convert::TryFromBorrowedObject, PyObject, PyResult, VirtualMachine};

vm/src/function/number.rs

@@ -1,5 +1,8 @@
+use super::argument::OptionalArg;
 use crate::{builtins::PyIntRef, AsObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine};
+use num_bigint::BigInt;
 use num_complex::Complex64;
+use num_traits::PrimInt;
 use std::ops::Deref;
 
 /// A Python complex-like object.
@@ -157,28 +160,39 @@ impl TryFromObject for ArgIndex {
 
 #[derive(Debug)]
 #[repr(transparent)]
-pub struct ArgSize {
-    value: isize,
+pub struct ArgPrimitiveIndex<T> {
+    pub value: T,
 }
 
-impl From<ArgSize> for isize {
-    fn from(arg: ArgSize) -> Self {
-        arg.value
+impl<T> OptionalArg<ArgPrimitiveIndex<T>> {
+    pub fn into_primitive(self) -> OptionalArg<T> {
+        self.map(|x| x.value)
     }
 }
 
-impl Deref for ArgSize {
-    type Target = isize;
+impl<T> Deref for ArgPrimitiveIndex<T> {
+    type Target = T;
 
     fn deref(&self) -> &Self::Target {
         &self.value
     }
 }
 
-impl TryFromObject for ArgSize {
+impl<T> TryFromObject for ArgPrimitiveIndex<T>
+where
+    T: PrimInt + for<'a> TryFrom<&'a BigInt>,
+{
     fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
         Ok(Self {
             value: obj.try_index(vm)?.try_to_primitive(vm)?,
         })
     }
 }
+
+pub type ArgSize = ArgPrimitiveIndex<isize>;
+
+impl From<ArgSize> for isize {
+    fn from(arg: ArgSize) -> Self {
+        arg.value
+    }
+}