Skip to content

Commit 3eafcbe

Browse files
committed
Add decorator functionality
1 parent 133d72f commit 3eafcbe

7 files changed

Lines changed: 190 additions & 50 deletions

File tree

parser/src/ast.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,13 +96,15 @@ pub enum Statement {
9696
name: String,
9797
body: Vec<LocatedStatement>,
9898
args: Vec<(String, Option<Expression>)>,
99+
decorator_list: Vec<Expression>,
99100
// TODO: docstring: String,
100101
},
101102
FunctionDef {
102103
name: String,
103104
args: Vec<(String, Option<Expression>)>,
104105
// docstring: String,
105106
body: Vec<LocatedStatement>,
107+
decorator_list: Vec<Expression>,
106108
},
107109
}
108110

@@ -132,6 +134,9 @@ pub enum Expression {
132134
op: UnaryOperator,
133135
a: Box<Expression>,
134136
},
137+
Yield {
138+
expression: Option<Box<Expression>>,
139+
},
135140
Compare {
136141
a: Box<Expression>,
137142
op: Comparison,

parser/src/parser.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,7 @@ mod tests {
307307
location: ast::Location::new(3, 3),
308308
node: ast::Statement::Pass,
309309
}],
310+
decorator_list: vec![],
310311
}
311312
},
312313
ast::LocatedStatement {
@@ -326,9 +327,11 @@ mod tests {
326327
location: ast::Location::new(5, 3),
327328
node: ast::Statement::Pass,
328329
}],
330+
decorator_list: vec![],
329331
}
330332
}
331333
],
334+
decorator_list: vec![],
332335
}
333336
})
334337
)

parser/src/python.lalrpop

Lines changed: 82 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,12 @@ FlowStatement: ast::LocatedStatement = {
137137
node: ast::Statement::Raise { expression: t },
138138
}
139139
},
140-
// yield
140+
<loc:@L> <y:YieldExpr> => {
141+
ast::LocatedStatement {
142+
location: loc,
143+
node: ast::Statement::Expression { expression: y },
144+
}
145+
},
141146
};
142147

143148
ImportStatement: ast::LocatedStatement = {
@@ -157,7 +162,7 @@ ImportStatement: ast::LocatedStatement = {
157162
},
158163
}
159164
},
160-
<loc:@L> "from" <n:DottedName> "import" <i: Comma<ImportPart<Identifier>>> => {
165+
<loc:@L> "from" <n:ImportFromLocation> "import" <i: ImportAsNames> => {
161166
ast::LocatedStatement {
162167
location: loc,
163168
node: ast::Statement::Import {
@@ -175,13 +180,50 @@ ImportStatement: ast::LocatedStatement = {
175180
},
176181
};
177182

183+
ImportFromLocation: String = {
184+
<dots: "."*> <name:DottedName> => {
185+
let mut r = "".to_string();
186+
for _dot in dots {
187+
r.push_str(".");
188+
}
189+
r.push_str(&name);
190+
r
191+
},
192+
<dots: "."+> => {
193+
let mut r = "".to_string();
194+
for _dot in dots {
195+
r.push_str(".");
196+
}
197+
r
198+
},
199+
};
200+
201+
ImportAsNames: Vec<(String, Option<String>)> = {
202+
<i:Comma<ImportPart<Identifier>>> => i,
203+
"(" <i:Comma<ImportPart<Identifier>>> ")" => i,
204+
"*" => {
205+
// Star import all
206+
vec![("*".to_string(), None)]
207+
},
208+
};
209+
210+
178211
#[inline]
179212
ImportPart<I>: (String, Option<String>) = {
180213
<i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
181214
};
182215

216+
// A name like abc or abc.def.ghi
183217
DottedName: String = {
184218
<n:name> => n,
219+
<n:name> <n2: ("." Identifier)+> => {
220+
let mut r = n.to_string();
221+
for x in n2 {
222+
r.push_str(".");
223+
r.push_str(&x.1);
224+
}
225+
r
226+
},
185227
};
186228

187229
AssertStatement: ast::LocatedStatement = {
@@ -322,10 +364,15 @@ WithItem: ast::WithItem = {
322364
};
323365

324366
FuncDef: ast::LocatedStatement = {
325-
<loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
367+
<d:Decorator*> <loc:@L> "def" <i:Identifier> <a:Parameters> ":" <s:Suite> => {
326368
ast::LocatedStatement {
327369
location: loc,
328-
node: ast::Statement::FunctionDef { name: i, args: a, body: s }
370+
node: ast::Statement::FunctionDef {
371+
name: i,
372+
args: a,
373+
body: s,
374+
decorator_list: d,
375+
}
329376
}
330377
},
331378
};
@@ -346,28 +393,51 @@ Parameter: (String, Option<ast::Expression>) = {
346393
};
347394

348395
ClassDef: ast::LocatedStatement = {
349-
<loc:@L> "class" <n:Identifier> <a:Parameters?> ":" <s:Suite> => {
396+
<d:Decorator*> <loc:@L> "class" <n:Identifier> <a:Parameters?> ":" <s:Suite> => {
350397
ast::LocatedStatement {
351398
location: loc,
352399
node: ast::Statement::ClassDef {
353400
name: n,
354401
args: a.unwrap_or(vec![]),
355-
body: s
402+
body: s,
403+
decorator_list: d,
356404
},
357405
}
358406
},
359407
};
360408

409+
// Decorators:
410+
Decorator: ast::Expression = {
411+
"@" <n:DottedName> <a: ("(" FunctionArguments ")")?> "\n" => {
412+
let name = ast::Expression::Identifier { name: n };
413+
match a {
414+
Some((_, args, _)) => ast::Expression::Call {
415+
function: Box::new(name),
416+
args: args,
417+
},
418+
None => name,
419+
}
420+
},
421+
};
422+
423+
YieldExpr: ast::Expression = {
424+
"yield" => {
425+
ast::Expression::Yield {
426+
expression: None,
427+
}
428+
}
429+
};
430+
361431
Test: ast::Expression = {
362432
<e:OrTest> => e,
363433
<e:LambdaDef> => e,
364434
};
365435

366436
LambdaDef: ast::Expression = {
367-
"lambda" <p:TypedArgsList> ":" <b:Expression> =>
368-
ast::Expression::Lambda {
369-
args:p,
370-
body:Box::new(b)
437+
"lambda" <p:TypedArgsList> ":" <b:Expression> =>
438+
ast::Expression::Lambda {
439+
args:p,
440+
body:Box::new(b)
371441
}
372442
}
373443

@@ -631,6 +701,7 @@ extern {
631701
"is" => lexer::Tok::Is,
632702
"import" => lexer::Tok::Import,
633703
"from" => lexer::Tok::From,
704+
"lambda" => lexer::Tok::Lambda,
634705
"not" => lexer::Tok::Not,
635706
"or" => lexer::Tok::Or,
636707
"pass" => lexer::Tok::Pass,
@@ -639,7 +710,7 @@ extern {
639710
"try" => lexer::Tok::Try,
640711
"while" => lexer::Tok::While,
641712
"with" => lexer::Tok::With,
642-
"lambda" => lexer::Tok::Lambda,
713+
"yield" => lexer::Tok::Yield,
643714
"True" => lexer::Tok::True,
644715
"False" => lexer::Tok::False,
645716
"None" => lexer::Tok::None,

src/main.rs

Lines changed: 33 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use clap::{App, Arg};
1111
use rustpython_parser::parser;
1212
use rustpython_vm::obj::objstr;
1313
use rustpython_vm::print_exception;
14-
use rustpython_vm::pyobject::PyObjectRef;
14+
use rustpython_vm::pyobject::{PyObjectRef, PyResult};
1515
use rustpython_vm::VirtualMachine;
1616
use rustpython_vm::{compile, import};
1717
use std::io;
@@ -45,67 +45,62 @@ fn main() {
4545
)
4646
.get_matches();
4747

48+
// Construct vm:
49+
let mut vm = VirtualMachine::new();
50+
4851
// Figure out if a -c option was given:
49-
if let Some(command) = matches.value_of("c") {
50-
run_command(&mut command.to_string());
52+
let result = if let Some(command) = matches.value_of("c") {
53+
run_command(&mut vm, command.to_string())
5154
} else if let Some(module) = matches.value_of("m") {
52-
run_module(module);
55+
run_module(&mut vm, module)
5356
} else {
5457
// Figure out if a script was passed:
5558
match matches.value_of("script") {
56-
None => run_shell(),
57-
Some(filename) => run_script(&filename.to_string()),
59+
None => run_shell(&mut vm),
60+
Some(filename) => run_script(&mut vm, &filename.to_string()),
5861
}
59-
}
62+
};
63+
64+
// See if any exception leaked out:
65+
handle_exception(&mut vm, result);
6066
}
6167

62-
fn _run_string(source: &str, source_path: Option<String>) {
63-
let mut vm = VirtualMachine::new();
64-
let code_obj = compile::compile(
65-
&mut vm,
66-
&source.to_string(),
67-
compile::Mode::Exec,
68-
source_path,
69-
)
70-
.unwrap();
68+
fn _run_string(vm: &mut VirtualMachine, source: &str, source_path: Option<String>) -> PyResult {
69+
let code_obj = compile::compile(vm, &source.to_string(), compile::Mode::Exec, source_path)?;
7170
debug!("Code object: {:?}", code_obj.borrow());
7271
let builtins = vm.get_builtin_scope();
7372
let vars = vm.context().new_scope(Some(builtins)); // Keep track of local variables
74-
match vm.run_code_obj(code_obj, vars) {
73+
vm.run_code_obj(code_obj, vars)
74+
}
75+
76+
fn handle_exception(vm: &mut VirtualMachine, result: PyResult) {
77+
match result {
7578
Ok(_value) => {}
76-
Err(exc) => {
77-
// println!("X: {:?}", exc.get_attr("__traceback__"));
78-
print_exception(&mut vm, &exc);
79-
panic!("Exception: {:?}", exc);
79+
Err(err) => {
80+
print_exception(vm, &err);
8081
}
8182
}
8283
}
8384

84-
fn run_command(source: &mut String) {
85+
fn run_command(vm: &mut VirtualMachine, mut source: String) -> PyResult {
8586
debug!("Running command {}", source);
8687

8788
// This works around https://github.com/RustPython/RustPython/issues/17
8889
source.push_str("\n");
89-
_run_string(source, None)
90+
_run_string(vm, &source, None)
9091
}
9192

92-
fn run_module(module: &str) {
93+
fn run_module(vm: &mut VirtualMachine, module: &str) -> PyResult {
9394
debug!("Running module {}", module);
94-
let mut vm = VirtualMachine::new();
95-
match import::import_module(&mut vm, module) {
96-
Ok(_value) => {}
97-
Err(err) => {
98-
print_exception(&mut vm, &err);
99-
}
100-
}
95+
import::import_module(vm, module)
10196
}
10297

103-
fn run_script(script_file: &str) {
98+
fn run_script(vm: &mut VirtualMachine, script_file: &str) -> PyResult {
10499
debug!("Running file {}", script_file);
105100
// Parse an ast from it:
106101
let filepath = Path::new(script_file);
107102
match parser::read_file(filepath) {
108-
Ok(source) => _run_string(&source, Some(filepath.to_str().unwrap().to_string())),
103+
Ok(source) => _run_string(vm, &source, Some(filepath.to_str().unwrap().to_string())),
109104
Err(msg) => {
110105
error!("Parsing went horribly wrong: {}", msg);
111106
std::process::exit(1);
@@ -141,12 +136,11 @@ fn shell_exec(vm: &mut VirtualMachine, source: &str, scope: PyObjectRef) -> bool
141136
true
142137
}
143138

144-
fn run_shell() {
139+
fn run_shell(vm: &mut VirtualMachine) -> PyResult {
145140
println!(
146141
"Welcome to the magnificent Rust Python {} interpreter",
147142
crate_version!()
148143
);
149-
let mut vm = VirtualMachine::new();
150144
let builtins = vm.get_builtin_scope();
151145
let vars = vm.context().new_scope(Some(builtins)); // Keep track of local variables
152146

@@ -167,7 +161,7 @@ fn run_shell() {
167161
}
168162
Ok(_) => {
169163
debug!("You entered {:?}", input);
170-
if shell_exec(&mut vm, &input, vars.clone()) {
164+
if shell_exec(vm, &input, vars.clone()) {
171165
// Line was complete.
172166
input = String::new();
173167
} else {
@@ -186,7 +180,7 @@ fn run_shell() {
186180
.trim_right_matches(|c| c == '\r' || c == '\n')
187181
.to_string();
188182
if line.len() == 0 {
189-
if shell_exec(&mut vm, &input, vars.clone()) {
183+
if shell_exec(vm, &input, vars.clone()) {
190184
input = String::new();
191185
break;
192186
}
@@ -203,4 +197,6 @@ fn run_shell() {
203197
Err(msg) => panic!("Error: {:?}", msg),
204198
};
205199
}
200+
201+
Ok(vm.get_none())
206202
}

tests/snippets/decorators.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
2+
def logged(f):
3+
def wrapper(a, b):
4+
print('Calling function', f)
5+
return f(a, b + 1)
6+
return wrapper
7+
8+
9+
@logged
10+
def add(a, b):
11+
return a + b
12+
13+
c = add(10, 3)
14+
15+
assert c == 14
16+

0 commit comments

Comments
 (0)