rustpython_vm/function/
argument.rs

1use crate::{
2    builtins::{PyBaseExceptionRef, PyTupleRef, PyTypeRef},
3    convert::ToPyObject,
4    object::{Traverse, TraverseFn},
5    AsObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
6};
7use indexmap::IndexMap;
8use itertools::Itertools;
9use std::ops::RangeInclusive;
10
11pub trait IntoFuncArgs: Sized {
12    fn into_args(self, vm: &VirtualMachine) -> FuncArgs;
13    fn into_method_args(self, obj: PyObjectRef, vm: &VirtualMachine) -> FuncArgs {
14        let mut args = self.into_args(vm);
15        args.prepend_arg(obj);
16        args
17    }
18}
19
20impl<T> IntoFuncArgs for T
21where
22    T: Into<FuncArgs>,
23{
24    fn into_args(self, _vm: &VirtualMachine) -> FuncArgs {
25        self.into()
26    }
27}
28
29// A tuple of values that each implement `ToPyObject` represents a sequence of
30// arguments that can be bound and passed to a built-in function.
31macro_rules! into_func_args_from_tuple {
32    ($(($n:tt, $T:ident)),*) => {
33        impl<$($T,)*> IntoFuncArgs for ($($T,)*)
34        where
35            $($T: ToPyObject,)*
36        {
37            #[inline]
38            fn into_args(self, vm: &VirtualMachine) -> FuncArgs {
39                let ($($n,)*) = self;
40                PosArgs::new(vec![$($n.to_pyobject(vm),)*]).into()
41            }
42
43            #[inline]
44            fn into_method_args(self, obj: PyObjectRef, vm: &VirtualMachine) -> FuncArgs {
45                let ($($n,)*) = self;
46                PosArgs::new(vec![obj, $($n.to_pyobject(vm),)*]).into()
47            }
48        }
49    };
50}
51
52into_func_args_from_tuple!((v1, T1));
53into_func_args_from_tuple!((v1, T1), (v2, T2));
54into_func_args_from_tuple!((v1, T1), (v2, T2), (v3, T3));
55into_func_args_from_tuple!((v1, T1), (v2, T2), (v3, T3), (v4, T4));
56into_func_args_from_tuple!((v1, T1), (v2, T2), (v3, T3), (v4, T4), (v5, T5));
57
58/// The `FuncArgs` struct is one of the most used structs then creating
59/// a rust function that can be called from python. It holds both positional
60/// arguments, as well as keyword arguments passed to the function.
61#[derive(Debug, Default, Clone, Traverse)]
62pub struct FuncArgs {
63    pub args: Vec<PyObjectRef>,
64    // sorted map, according to https://www.python.org/dev/peps/pep-0468/
65    pub kwargs: IndexMap<String, PyObjectRef>,
66}
67
68unsafe impl Traverse for IndexMap<String, PyObjectRef> {
69    fn traverse(&self, tracer_fn: &mut TraverseFn) {
70        self.values().for_each(|v| v.traverse(tracer_fn));
71    }
72}
73
74/// Conversion from vector of python objects to function arguments.
75impl<A> From<A> for FuncArgs
76where
77    A: Into<PosArgs>,
78{
79    fn from(args: A) -> Self {
80        FuncArgs {
81            args: args.into().into_vec(),
82            kwargs: IndexMap::new(),
83        }
84    }
85}
86
87impl From<KwArgs> for FuncArgs {
88    fn from(kwargs: KwArgs) -> Self {
89        FuncArgs {
90            args: Vec::new(),
91            kwargs: kwargs.0,
92        }
93    }
94}
95
96impl FromArgs for FuncArgs {
97    fn from_args(_vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
98        Ok(std::mem::take(args))
99    }
100}
101
102impl FuncArgs {
103    pub fn new<A, K>(args: A, kwargs: K) -> Self
104    where
105        A: Into<PosArgs>,
106        K: Into<KwArgs>,
107    {
108        let PosArgs(args) = args.into();
109        let KwArgs(kwargs) = kwargs.into();
110        Self { args, kwargs }
111    }
112
113    pub fn with_kwargs_names<A, KW>(mut args: A, kwarg_names: KW) -> Self
114    where
115        A: ExactSizeIterator<Item = PyObjectRef>,
116        KW: ExactSizeIterator<Item = String>,
117    {
118        // last `kwarg_names.len()` elements of args in order of appearance in the call signature
119        let total_argc = args.len();
120        let kwargc = kwarg_names.len();
121        let posargc = total_argc - kwargc;
122
123        let posargs = args.by_ref().take(posargc).collect();
124
125        let kwargs = kwarg_names.zip_eq(args).collect::<IndexMap<_, _>>();
126
127        FuncArgs {
128            args: posargs,
129            kwargs,
130        }
131    }
132
133    pub fn is_empty(&self) -> bool {
134        self.args.is_empty() && self.kwargs.is_empty()
135    }
136
137    pub fn prepend_arg(&mut self, item: PyObjectRef) {
138        self.args.reserve_exact(1);
139        self.args.insert(0, item)
140    }
141
142    pub fn shift(&mut self) -> PyObjectRef {
143        self.args.remove(0)
144    }
145
146    pub fn get_kwarg(&self, key: &str, default: PyObjectRef) -> PyObjectRef {
147        self.kwargs
148            .get(key)
149            .cloned()
150            .unwrap_or_else(|| default.clone())
151    }
152
153    pub fn get_optional_kwarg(&self, key: &str) -> Option<PyObjectRef> {
154        self.kwargs.get(key).cloned()
155    }
156
157    pub fn get_optional_kwarg_with_type(
158        &self,
159        key: &str,
160        ty: PyTypeRef,
161        vm: &VirtualMachine,
162    ) -> PyResult<Option<PyObjectRef>> {
163        match self.get_optional_kwarg(key) {
164            Some(kwarg) => {
165                if kwarg.fast_isinstance(&ty) {
166                    Ok(Some(kwarg))
167                } else {
168                    let expected_ty_name = &ty.name();
169                    let kwarg_class = kwarg.class();
170                    let actual_ty_name = &kwarg_class.name();
171                    Err(vm.new_type_error(format!(
172                        "argument of type {expected_ty_name} is required for named parameter `{key}` (got: {actual_ty_name})"
173                    )))
174                }
175            }
176            None => Ok(None),
177        }
178    }
179
180    pub fn take_positional(&mut self) -> Option<PyObjectRef> {
181        if self.args.is_empty() {
182            None
183        } else {
184            Some(self.args.remove(0))
185        }
186    }
187
188    pub fn take_positional_keyword(&mut self, name: &str) -> Option<PyObjectRef> {
189        self.take_positional().or_else(|| self.take_keyword(name))
190    }
191
192    pub fn take_keyword(&mut self, name: &str) -> Option<PyObjectRef> {
193        self.kwargs.swap_remove(name)
194    }
195
196    pub fn remaining_keywords(&mut self) -> impl Iterator<Item = (String, PyObjectRef)> + '_ {
197        self.kwargs.drain(..)
198    }
199
200    /// Binds these arguments to their respective values.
201    ///
202    /// If there is an insufficient number of arguments, there are leftover
203    /// arguments after performing the binding, or if an argument is not of
204    /// the expected type, a TypeError is raised.
205    ///
206    /// If the given `FromArgs` includes any conversions, exceptions raised
207    /// during the conversion will halt the binding and return the error.
208    pub fn bind<T: FromArgs>(mut self, vm: &VirtualMachine) -> PyResult<T> {
209        let given_args = self.args.len();
210        let bound = T::from_args(vm, &mut self)
211            .map_err(|e| e.into_exception(T::arity(), given_args, vm))?;
212
213        if !self.args.is_empty() {
214            Err(vm.new_type_error(format!(
215                "Expected at most {} arguments ({} given)",
216                T::arity().end(),
217                given_args,
218            )))
219        } else if let Some(err) = self.check_kwargs_empty(vm) {
220            Err(err)
221        } else {
222            Ok(bound)
223        }
224    }
225
226    pub fn check_kwargs_empty(&self, vm: &VirtualMachine) -> Option<PyBaseExceptionRef> {
227        self.kwargs
228            .keys()
229            .next()
230            .map(|k| vm.new_type_error(format!("Unexpected keyword argument {k}")))
231    }
232}
233
234/// An error encountered while binding arguments to the parameters of a Python
235/// function call.
236pub enum ArgumentError {
237    /// The call provided fewer positional arguments than the function requires.
238    TooFewArgs,
239    /// The call provided more positional arguments than the function accepts.
240    TooManyArgs,
241    /// The function doesn't accept a keyword argument with the given name.
242    InvalidKeywordArgument(String),
243    /// The function require a keyword argument with the given name, but one wasn't provided
244    RequiredKeywordArgument(String),
245    /// An exception was raised while binding arguments to the function
246    /// parameters.
247    Exception(PyBaseExceptionRef),
248}
249
250impl From<PyBaseExceptionRef> for ArgumentError {
251    fn from(ex: PyBaseExceptionRef) -> Self {
252        ArgumentError::Exception(ex)
253    }
254}
255
256impl ArgumentError {
257    fn into_exception(
258        self,
259        arity: RangeInclusive<usize>,
260        num_given: usize,
261        vm: &VirtualMachine,
262    ) -> PyBaseExceptionRef {
263        match self {
264            ArgumentError::TooFewArgs => vm.new_type_error(format!(
265                "Expected at least {} arguments ({} given)",
266                arity.start(),
267                num_given
268            )),
269            ArgumentError::TooManyArgs => vm.new_type_error(format!(
270                "Expected at most {} arguments ({} given)",
271                arity.end(),
272                num_given
273            )),
274            ArgumentError::InvalidKeywordArgument(name) => {
275                vm.new_type_error(format!("{name} is an invalid keyword argument"))
276            }
277            ArgumentError::RequiredKeywordArgument(name) => {
278                vm.new_type_error(format!("Required keyqord only argument {name}"))
279            }
280            ArgumentError::Exception(ex) => ex,
281        }
282    }
283}
284
285/// Implemented by any type that can be accepted as a parameter to a built-in
286/// function.
287///
288pub trait FromArgs: Sized {
289    /// The range of positional arguments permitted by the function signature.
290    ///
291    /// Returns an empty range if not applicable.
292    fn arity() -> RangeInclusive<usize> {
293        0..=0
294    }
295
296    /// Extracts this item from the next argument(s).
297    fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError>;
298}
299
300pub trait FromArgOptional {
301    type Inner: TryFromObject;
302    fn from_inner(x: Self::Inner) -> Self;
303}
304impl<T: TryFromObject> FromArgOptional for OptionalArg<T> {
305    type Inner = T;
306    fn from_inner(x: T) -> Self {
307        Self::Present(x)
308    }
309}
310impl<T: TryFromObject> FromArgOptional for T {
311    type Inner = Self;
312    fn from_inner(x: Self) -> Self {
313        x
314    }
315}
316
317/// A map of keyword arguments to their values.
318///
319/// A built-in function with a `KwArgs` parameter is analogous to a Python
320/// function with `**kwargs`. All remaining keyword arguments are extracted
321/// (and hence the function will permit an arbitrary number of them).
322///
323/// `KwArgs` optionally accepts a generic type parameter to allow type checks
324/// or conversions of each argument.
325///
326/// Note:
327///
328/// KwArgs is only for functions that accept arbitrary keyword arguments. For
329/// functions that accept only *specific* named arguments, a rust struct with
330/// an appropriate FromArgs implementation must be created.
331#[derive(Clone)]
332pub struct KwArgs<T = PyObjectRef>(IndexMap<String, T>);
333
334unsafe impl<T> Traverse for KwArgs<T>
335where
336    T: Traverse,
337{
338    fn traverse(&self, tracer_fn: &mut TraverseFn) {
339        self.0.iter().map(|(_, v)| v.traverse(tracer_fn)).count();
340    }
341}
342
343impl<T> KwArgs<T> {
344    pub fn new(map: IndexMap<String, T>) -> Self {
345        KwArgs(map)
346    }
347
348    pub fn pop_kwarg(&mut self, name: &str) -> Option<T> {
349        self.0.swap_remove(name)
350    }
351
352    pub fn is_empty(self) -> bool {
353        self.0.is_empty()
354    }
355}
356impl<T> FromIterator<(String, T)> for KwArgs<T> {
357    fn from_iter<I: IntoIterator<Item = (String, T)>>(iter: I) -> Self {
358        KwArgs(iter.into_iter().collect())
359    }
360}
361impl<T> Default for KwArgs<T> {
362    fn default() -> Self {
363        KwArgs(IndexMap::new())
364    }
365}
366
367impl<T> FromArgs for KwArgs<T>
368where
369    T: TryFromObject,
370{
371    fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
372        let mut kwargs = IndexMap::new();
373        for (name, value) in args.remaining_keywords() {
374            kwargs.insert(name, value.try_into_value(vm)?);
375        }
376        Ok(KwArgs(kwargs))
377    }
378}
379
380impl<T> IntoIterator for KwArgs<T> {
381    type Item = (String, T);
382    type IntoIter = indexmap::map::IntoIter<String, T>;
383
384    fn into_iter(self) -> Self::IntoIter {
385        self.0.into_iter()
386    }
387}
388
389/// A list of positional argument values.
390///
391/// A built-in function with a `PosArgs` parameter is analogous to a Python
392/// function with `*args`. All remaining positional arguments are extracted
393/// (and hence the function will permit an arbitrary number of them).
394///
395/// `PosArgs` optionally accepts a generic type parameter to allow type checks
396/// or conversions of each argument.
397#[derive(Clone)]
398pub struct PosArgs<T = PyObjectRef>(Vec<T>);
399
400unsafe impl<T> Traverse for PosArgs<T>
401where
402    T: Traverse,
403{
404    fn traverse(&self, tracer_fn: &mut TraverseFn) {
405        self.0.traverse(tracer_fn)
406    }
407}
408
409impl<T> PosArgs<T> {
410    pub fn new(args: Vec<T>) -> Self {
411        Self(args)
412    }
413
414    pub fn into_vec(self) -> Vec<T> {
415        self.0
416    }
417
418    pub fn iter(&self) -> std::slice::Iter<T> {
419        self.0.iter()
420    }
421}
422
423impl<T> From<Vec<T>> for PosArgs<T> {
424    fn from(v: Vec<T>) -> Self {
425        Self(v)
426    }
427}
428
429impl From<()> for PosArgs<PyObjectRef> {
430    fn from(_args: ()) -> Self {
431        Self(Vec::new())
432    }
433}
434
435impl<T> AsRef<[T]> for PosArgs<T> {
436    fn as_ref(&self) -> &[T] {
437        &self.0
438    }
439}
440
441impl<T: PyPayload> PosArgs<PyRef<T>> {
442    pub fn into_tuple(self, vm: &VirtualMachine) -> PyTupleRef {
443        vm.ctx
444            .new_tuple(self.0.into_iter().map(Into::into).collect())
445    }
446}
447
448impl<T> FromArgs for PosArgs<T>
449where
450    T: TryFromObject,
451{
452    fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
453        let mut varargs = Vec::new();
454        while let Some(value) = args.take_positional() {
455            varargs.push(value.try_into_value(vm)?);
456        }
457        Ok(PosArgs(varargs))
458    }
459}
460
461impl<T> IntoIterator for PosArgs<T> {
462    type Item = T;
463    type IntoIter = std::vec::IntoIter<T>;
464
465    fn into_iter(self) -> Self::IntoIter {
466        self.0.into_iter()
467    }
468}
469
470impl<T> FromArgs for T
471where
472    T: TryFromObject,
473{
474    fn arity() -> RangeInclusive<usize> {
475        1..=1
476    }
477
478    fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
479        let value = args.take_positional().ok_or(ArgumentError::TooFewArgs)?;
480        Ok(value.try_into_value(vm)?)
481    }
482}
483
484/// An argument that may or may not be provided by the caller.
485///
486/// This style of argument is not possible in pure Python.
487#[derive(Debug, result_like::OptionLike, is_macro::Is)]
488pub enum OptionalArg<T = PyObjectRef> {
489    Present(T),
490    Missing,
491}
492
493unsafe impl<T> Traverse for OptionalArg<T>
494where
495    T: Traverse,
496{
497    fn traverse(&self, tracer_fn: &mut TraverseFn) {
498        match self {
499            OptionalArg::Present(ref o) => o.traverse(tracer_fn),
500            OptionalArg::Missing => (),
501        }
502    }
503}
504
505impl OptionalArg<PyObjectRef> {
506    pub fn unwrap_or_none(self, vm: &VirtualMachine) -> PyObjectRef {
507        self.unwrap_or_else(|| vm.ctx.none())
508    }
509}
510
511pub type OptionalOption<T = PyObjectRef> = OptionalArg<Option<T>>;
512
513impl<T> OptionalOption<T> {
514    #[inline]
515    pub fn flatten(self) -> Option<T> {
516        self.into_option().flatten()
517    }
518}
519
520impl<T> FromArgs for OptionalArg<T>
521where
522    T: TryFromObject,
523{
524    fn arity() -> RangeInclusive<usize> {
525        0..=1
526    }
527
528    fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
529        let r = if let Some(value) = args.take_positional() {
530            OptionalArg::Present(value.try_into_value(vm)?)
531        } else {
532            OptionalArg::Missing
533        };
534        Ok(r)
535    }
536}
537
538// For functions that accept no arguments. Implemented explicitly instead of via
539// macro below to avoid unused warnings.
540impl FromArgs for () {
541    fn from_args(_vm: &VirtualMachine, _args: &mut FuncArgs) -> Result<Self, ArgumentError> {
542        Ok(())
543    }
544}
545
546// A tuple of types that each implement `FromArgs` represents a sequence of
547// arguments that can be bound and passed to a built-in function.
548//
549// Technically, a tuple can contain tuples, which can contain tuples, and so on,
550// so this actually represents a tree of values to be bound from arguments, but
551// in practice this is only used for the top-level parameters.
552macro_rules! tuple_from_py_func_args {
553    ($($T:ident),+) => {
554        impl<$($T),+> FromArgs for ($($T,)+)
555        where
556            $($T: FromArgs),+
557        {
558            fn arity() -> RangeInclusive<usize> {
559                let mut min = 0;
560                let mut max = 0;
561                $(
562                    let (start, end) = $T::arity().into_inner();
563                    min += start;
564                    max += end;
565                )+
566                min..=max
567            }
568
569            fn from_args(vm: &VirtualMachine, args: &mut FuncArgs) -> Result<Self, ArgumentError> {
570                Ok(($($T::from_args(vm, args)?,)+))
571            }
572        }
573    };
574}
575
576// Implement `FromArgs` for up to 7-tuples, allowing built-in functions to bind
577// up to 7 top-level parameters (note that `PosArgs`, `KwArgs`, nested tuples, etc.
578// count as 1, so this should actually be more than enough).
579tuple_from_py_func_args!(A);
580tuple_from_py_func_args!(A, B);
581tuple_from_py_func_args!(A, B, C);
582tuple_from_py_func_args!(A, B, C, D);
583tuple_from_py_func_args!(A, B, C, D, E);
584tuple_from_py_func_args!(A, B, C, D, E, F);
585tuple_from_py_func_args!(A, B, C, D, E, F, G);
586tuple_from_py_func_args!(A, B, C, D, E, F, G, H);
rustpython_vm/function/argument.rs

rustpython_vm/function/
argument.rs
