Skip to content

Commit 3ca387b

Browse files
skinnyBatskinnyBat
committed
Enhance FromArgs derive proc macro to allow positional, and positional
or keyword arguments.
1 parent 17b816f commit 3ca387b

File tree

5 files changed

+188
-42
lines changed

5 files changed

+188
-42
lines changed

‎derive/src/lib.rs

Lines changed: 44 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,31 +3,64 @@ extern crate proc_macro;
33
use proc_macro::TokenStream;
44
use proc_macro2::TokenStream as TokenStream2;
55
use quote::quote;
6-
use syn::{Data, DeriveInput, Fields};
6+
use syn::{Data, DeriveInput, Field, Fields};
77

8-
#[proc_macro_derive(FromArgs)]
8+
#[proc_macro_derive(FromArgs, attributes(positional, keyword))]
99
pub fn derive_from_args(input: TokenStream) -> TokenStream {
1010
let ast: DeriveInput = syn::parse(input).unwrap();
1111

1212
let gen = impl_from_args(&ast);
1313
gen.to_string().parse().unwrap()
1414
}
1515

16+
enum ArgType {
17+
Positional,
18+
PositionalKeyword,
19+
Keyword,
20+
}
21+
22+
fn generate_field(field: &Field) -> TokenStream2 {
23+
let arg_type = if let Some(attr) = field.attrs.first() {
24+
if attr.path.is_ident("positional") {
25+
ArgType::Positional
26+
} else if attr.path.is_ident("keyword") {
27+
ArgType::Keyword
28+
} else {
29+
panic!("Unrecognised attribute")
30+
}
31+
} else {
32+
ArgType::PositionalKeyword
33+
};
34+
35+
let name = &field.ident;
36+
match arg_type {
37+
ArgType::Positional => {
38+
quote! {
39+
#name: args.take_positional(vm)?,
40+
}
41+
}
42+
ArgType::PositionalKeyword => {
43+
quote! {
44+
#name: args.take_positional_keyword(vm, stringify!(#name))?,
45+
}
46+
}
47+
ArgType::Keyword => {
48+
quote! {
49+
#name: args.take_keyword(vm, stringify!(#name))?,
50+
}
51+
}
52+
}
53+
}
54+
1655
fn impl_from_args(input: &DeriveInput) -> TokenStream2 {
1756
// FIXME: This references types using `crate` instead of `rustpython_vm`
1857
// so that it can be used in the latter. How can we support both?
58+
// Can use extern crate self as rustpython_vm; once in stable.
59+
// https://github.com/rust-lang/rust/issues/56409
1960
let fields = match input.data {
2061
Data::Struct(ref data) => {
2162
match data.fields {
22-
Fields::Named(ref fields) => fields.named.iter().map(|field| {
23-
let name = &field.ident;
24-
quote! {
25-
#name: crate::pyobject::TryFromObject::try_from_object(
26-
vm,
27-
args.take_keyword(stringify!(#name)).unwrap_or_else(|| vm.ctx.none())
28-
)?,
29-
}
30-
}),
63+
Fields::Named(ref fields) => fields.named.iter().map(generate_field),
3164
Fields::Unnamed(_) | Fields::Unit => unimplemented!(), // TODO: better error message
3265
}
3366
}

‎tests/snippets/ints.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from testutils import assert_raises
1+
from testutils import assert_raises, assertRaises
22

33
# int to int comparisons
44

@@ -58,3 +58,19 @@
5858
assert (2).__rmul__(1.0) == NotImplemented
5959
assert (2).__truediv__(1.0) == NotImplemented
6060
assert (2).__rtruediv__(1.0) == NotImplemented
61+
62+
63+
assert int() == 0
64+
assert int("101", 2) == 5
65+
assert int("101", base=2) == 5
66+
assert int(1) == 1
67+
68+
with assertRaises(TypeError):
69+
int(base=2)
70+
71+
with assertRaises(TypeError):
72+
int(1, base=2)
73+
74+
with assertRaises(TypeError):
75+
# check that first parameter is truly positional only
76+
int(val_options=1)

‎vm/src/builtins.rs

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -559,11 +559,22 @@ fn builtin_pow(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
559559
}
560560
}
561561

562+
// Idea: Should we have a 'default' attribute, so we don't have to use OptionalArg's in this case
563+
//#[derive(Debug, FromArgs)]
564+
//pub struct PrintOptions {
565+
// #[keyword] #[default(None)] sep: Option<PyStringRef>,
566+
// #[keyword] #[default(None)] end: Option<PyStringRef>,
567+
// #[keyword] #[default(flush)] flush: bool,
568+
//}
569+
562570
#[derive(Debug, FromArgs)]
563571
pub struct PrintOptions {
564-
sep: Option<PyStringRef>,
565-
end: Option<PyStringRef>,
566-
flush: bool,
572+
#[keyword]
573+
sep: OptionalArg<Option<PyStringRef>>,
574+
#[keyword]
575+
end: OptionalArg<Option<PyStringRef>>,
576+
#[keyword]
577+
flush: OptionalArg<bool>,
567578
}
568579

569580
pub fn builtin_print(objects: Args, options: PrintOptions, vm: &VirtualMachine) -> PyResult<()> {
@@ -573,7 +584,7 @@ pub fn builtin_print(objects: Args, options: PrintOptions, vm: &VirtualMachine)
573584
for object in objects {
574585
if first {
575586
first = false;
576-
} else if let Some(ref sep) = options.sep {
587+
} else if let OptionalArg::Present(Some(ref sep)) = options.sep {
577588
write!(stdout_lock, "{}", sep.value).unwrap();
578589
} else {
579590
write!(stdout_lock, " ").unwrap();
@@ -582,13 +593,13 @@ pub fn builtin_print(objects: Args, options: PrintOptions, vm: &VirtualMachine)
582593
write!(stdout_lock, "{}", s).unwrap();
583594
}
584595

585-
if let Some(end) = options.end {
596+
if let OptionalArg::Present(Some(end)) = options.end {
586597
write!(stdout_lock, "{}", end.value).unwrap();
587598
} else {
588599
writeln!(stdout_lock).unwrap();
589600
}
590601

591-
if options.flush {
602+
if options.flush.into_option().unwrap_or(false) {
592603
stdout_lock.flush().unwrap();
593604
}
594605

‎vm/src/function.rs

Lines changed: 78 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
use std::collections::HashMap;
22
use std::ops::RangeInclusive;
33

4-
use crate::obj::objtype;
4+
use crate::obj::objtype::{isinstance, PyClassRef};
55
use crate::pyobject::{IntoPyObject, PyObjectRef, PyResult, TryFromObject, TypeProtocol};
66
use crate::vm::VirtualMachine;
77

88
use self::OptionalArg::*;
9-
use crate::obj::objtype::PyClassRef;
109

1110
/// The `PyFuncArgs` struct is one of the most used structs then creating
1211
/// a rust function that can be called from python. It holds both positional
@@ -95,7 +94,7 @@ impl PyFuncArgs {
9594
) -> Result<Option<PyObjectRef>, PyObjectRef> {
9695
match self.get_optional_kwarg(key) {
9796
Some(kwarg) => {
98-
if objtype::isinstance(&kwarg, &ty) {
97+
if isinstance(&kwarg, &ty) {
9998
Ok(Some(kwarg))
10099
} else {
101100
let expected_ty_name = vm.to_pystr(&ty)?;
@@ -118,7 +117,7 @@ impl PyFuncArgs {
118117
}
119118
}
120119

121-
pub fn take_keyword(&mut self, name: &str) -> Option<PyObjectRef> {
120+
fn extract_keyword(&mut self, name: &str) -> Option<PyObjectRef> {
122121
// TODO: change kwarg representation so this scan isn't necessary
123122
if let Some(index) = self
124123
.kwargs
@@ -131,6 +130,49 @@ impl PyFuncArgs {
131130
}
132131
}
133132

133+
pub fn take_positional<H: ArgHandler>(
134+
&mut self,
135+
vm: &VirtualMachine,
136+
) -> Result<H, ArgumentError> {
137+
if let Some(arg) = self.next_positional() {
138+
H::from_arg(vm, arg).map_err(|err| ArgumentError::Exception(err))
139+
} else if let Some(default) = H::DEFAULT {
140+
Ok(default)
141+
} else {
142+
Err(ArgumentError::TooFewArgs)
143+
}
144+
}
145+
146+
pub fn take_positional_keyword<H: ArgHandler>(
147+
&mut self,
148+
vm: &VirtualMachine,
149+
name: &str,
150+
) -> Result<H, ArgumentError> {
151+
if let Some(arg) = self.next_positional() {
152+
H::from_arg(vm, arg).map_err(|err| ArgumentError::Exception(err))
153+
} else if let Some(arg) = self.extract_keyword(name) {
154+
H::from_arg(vm, arg).map_err(|err| ArgumentError::Exception(err))
155+
} else if let Some(default) = H::DEFAULT {
156+
Ok(default)
157+
} else {
158+
Err(ArgumentError::TooFewArgs)
159+
}
160+
}
161+
162+
pub fn take_keyword<H: ArgHandler>(
163+
&mut self,
164+
vm: &VirtualMachine,
165+
name: &str,
166+
) -> Result<H, ArgumentError> {
167+
if let Some(arg) = self.extract_keyword(name) {
168+
H::from_arg(vm, arg).map_err(|err| ArgumentError::Exception(err))
169+
} else if let Some(default) = H::DEFAULT {
170+
Ok(default)
171+
} else {
172+
Err(ArgumentError::RequiredKeywordArgument(name.to_string()))
173+
}
174+
}
175+
134176
pub fn remaining_keyword<'a>(&'a mut self) -> impl Iterator<Item = (String, PyObjectRef)> + 'a {
135177
self.kwargs.drain(..)
136178
}
@@ -164,6 +206,9 @@ impl PyFuncArgs {
164206
Err(ArgumentError::InvalidKeywordArgument(name)) => {
165207
return Err(vm.new_type_error(format!("{} is an invalid keyword argument", name)));
166208
}
209+
Err(ArgumentError::RequiredKeywordArgument(name)) => {
210+
return Err(vm.new_type_error(format!("Required keyqord only argument {}", name)));
211+
}
167212
Err(ArgumentError::Exception(ex)) => {
168213
return Err(ex);
169214
}
@@ -192,6 +237,8 @@ pub enum ArgumentError {
192237
TooManyArgs,
193238
/// The function doesn't accept a keyword argument with the given name.
194239
InvalidKeywordArgument(String),
240+
/// The function require a keyword argument with the given name, but one wasn't provided
241+
RequiredKeywordArgument(String),
195242
/// An exception was raised while binding arguments to the function
196243
/// parameters.
197244
Exception(PyObjectRef),
@@ -218,6 +265,32 @@ pub trait FromArgs: Sized {
218265
fn from_args(vm: &VirtualMachine, args: &mut PyFuncArgs) -> Result<Self, ArgumentError>;
219266
}
220267

268+
/// Handling behaviour when the argument is and isn't present, used to implement OptionalArg.
269+
pub trait ArgHandler: Sized {
270+
/// Default value that will be used if the argument isn't present, or None in which case the a
271+
/// appropriate error is returned.
272+
const DEFAULT: Option<Self>;
273+
274+
/// Converts the arg argument when it is present
275+
fn from_arg(vm: &VirtualMachine, object: PyObjectRef) -> PyResult<Self>;
276+
}
277+
278+
impl<T: TryFromObject> ArgHandler for OptionalArg<T> {
279+
const DEFAULT: Option<Self> = Some(Missing);
280+
281+
fn from_arg(vm: &VirtualMachine, object: PyObjectRef) -> PyResult<Self> {
282+
T::try_from_object(vm, object).map(|x| Present(x))
283+
}
284+
}
285+
286+
impl<T: TryFromObject> ArgHandler for T {
287+
const DEFAULT: Option<Self> = None;
288+
289+
fn from_arg(vm: &VirtualMachine, object: PyObjectRef) -> PyResult<Self> {
290+
T::try_from_object(vm, object)
291+
}
292+
}
293+
221294
/// A map of keyword arguments to their values.
222295
///
223296
/// A built-in function with a `KwArgs` parameter is analagous to a Python
@@ -308,6 +381,7 @@ where
308381
/// An argument that may or may not be provided by the caller.
309382
///
310383
/// This style of argument is not possible in pure Python.
384+
#[derive(Debug)]
311385
pub enum OptionalArg<T> {
312386
Present(T),
313387
Missing,

‎vm/src/obj/objint.rs

Lines changed: 32 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use num_integer::Integer;
66
use num_traits::{Pow, Signed, ToPrimitive, Zero};
77

88
use crate::format::FormatSpec;
9-
use crate::function::{OptionalArg, PyFuncArgs};
9+
use crate::function::OptionalArg;
1010
use crate::pyobject::{
1111
IntoPyObject, PyContext, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
1212
};
@@ -380,25 +380,37 @@ impl PyIntRef {
380380
}
381381
}
382382

383-
fn int_new(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
384-
arg_check!(
385-
vm,
386-
args,
387-
required = [(cls, None)],
388-
optional = [(val_option, None)]
389-
);
390-
391-
let base = match args.get_optional_kwarg("base") {
392-
Some(argument) => get_value(&argument).to_u32().unwrap(),
393-
None => 10,
394-
};
395-
let val = match val_option {
396-
Some(val) => to_int(vm, val, base)?,
397-
None => Zero::zero(),
398-
};
399-
Ok(PyInt::new(val)
400-
.into_ref_with_type(vm, cls.clone().downcast().unwrap())?
401-
.into_object())
383+
#[derive(FromArgs)]
384+
struct IntOptions {
385+
#[positional]
386+
val_options: OptionalArg<PyObjectRef>,
387+
base: OptionalArg<u32>,
388+
}
389+
390+
impl IntOptions {
391+
fn get_int_value(self, vm: &VirtualMachine) -> PyResult<BigInt> {
392+
if let OptionalArg::Present(val) = self.val_options {
393+
let base = if let OptionalArg::Present(base) = self.base {
394+
if !objtype::isinstance(&val, &vm.ctx.str_type) {
395+
return Err(vm.new_type_error(
396+
"int() can't convert non-string with explicit base".to_string(),
397+
));
398+
}
399+
base
400+
} else {
401+
10
402+
};
403+
to_int(vm, &val, base)
404+
} else if let OptionalArg::Present(_) = self.base {
405+
Err(vm.new_type_error("int() missing string argument".to_string()))
406+
} else {
407+
Ok(Zero::zero())
408+
}
409+
}
410+
}
411+
412+
fn int_new(cls: PyClassRef, options: IntOptions, vm: &VirtualMachine) -> PyResult<PyIntRef> {
413+
PyInt::new(options.get_int_value(vm)?).into_ref_with_type(vm, cls)
402414
}
403415

404416
// Casting function:

0 commit comments

Comments
 (0)