Skip to content

Commit 0741c36

Browse files
windelbouwmanwindelbouwman
committed
Simplify import AST in line with CPython.
1 parent a0abe08 commit 0741c36

File tree

7 files changed

+124
-130
lines changed

7 files changed

+124
-130
lines changed

‎bytecode/src/bytecode.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ pub enum Instruction {
8888
name: String,
8989
level: usize,
9090
},
91+
ImportFrom {
92+
name: String,
93+
},
9194
LoadName {
9295
name: String,
9396
scope: NameScope,
@@ -379,6 +382,7 @@ impl Instruction {
379382
level,
380383
} => w!(Import, name, format!("{:?}", symbols), level),
381384
ImportStar { name, level } => w!(ImportStar, name, level),
385+
ImportFrom { name } => w!(ImportFrom, name),
382386
LoadName { name, scope } => w!(LoadName, name, format!("{:?}", scope)),
383387
StoreName { name, scope } => w!(StoreName, name, format!("{:?}", scope)),
384388
DeleteName { name } => w!(DeleteName, name),

‎compiler/src/compile.rs‎

Lines changed: 53 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -267,61 +267,63 @@ impl Compiler {
267267
self.set_source_location(&statement.location);
268268

269269
match &statement.node {
270-
ast::Statement::Import { import_parts } => {
271-
for ast::SingleImport {
272-
module,
273-
symbols,
274-
alias,
275-
level,
276-
} in import_parts
277-
{
278-
let level = *level;
279-
if let Some(alias) = alias {
280-
// import module as alias
281-
self.emit(Instruction::Import {
282-
name: module.clone(),
283-
symbols: vec![],
284-
level,
285-
});
286-
self.store_name(&alias);
287-
} else if symbols.is_empty() {
288-
// import module
289-
self.emit(Instruction::Import {
290-
name: module.clone(),
291-
symbols: vec![],
292-
level,
293-
});
294-
self.store_name(&module.clone());
270+
ast::Statement::Import { names } => {
271+
// import a, b, c as d
272+
for name in names {
273+
self.emit(Instruction::Import {
274+
name: name.symbol.clone(),
275+
symbols: vec![],
276+
level: 0,
277+
});
278+
279+
if let Some(alias) = &name.alias {
280+
self.store_name(alias);
295281
} else {
296-
let import_star = symbols
297-
.iter()
298-
.any(|import_symbol| import_symbol.symbol == "*");
299-
if import_star {
300-
// from module import *
301-
self.emit(Instruction::ImportStar {
302-
name: module.clone(),
303-
level,
304-
});
282+
self.store_name(&name.symbol);
283+
}
284+
}
285+
}
286+
ast::Statement::ImportFrom {
287+
level,
288+
module,
289+
names,
290+
} => {
291+
let import_star = names.iter().any(|n| n.symbol == "*");
292+
293+
if import_star {
294+
// from .... import *
295+
self.emit(Instruction::ImportStar {
296+
name: module.clone().unwrap(),
297+
level: *level,
298+
});
299+
} else {
300+
// from mod import a, b as c
301+
// First, determine the fromlist (for import lib):
302+
let from_list = names.iter().map(|n| n.symbol.clone()).collect();
303+
304+
// Load module once:
305+
self.emit(Instruction::Import {
306+
name: module.clone().unwrap(),
307+
symbols: from_list,
308+
level: *level,
309+
});
310+
311+
for name in names {
312+
// import symbol from module:
313+
self.emit(Instruction::ImportFrom {
314+
name: name.symbol.to_string(),
315+
});
316+
317+
// Store module under proper name:
318+
if let Some(alias) = &name.alias {
319+
self.store_name(alias);
305320
} else {
306-
// from module import symbol
307-
// from module import symbol as alias
308-
let (names, symbols_strings): (Vec<String>, Vec<String>) = symbols
309-
.iter()
310-
.map(|ast::ImportSymbol { symbol, alias }| {
311-
(
312-
alias.clone().unwrap_or_else(|| symbol.to_string()),
313-
symbol.to_string(),
314-
)
315-
})
316-
.unzip();
317-
self.emit(Instruction::Import {
318-
name: module.clone(),
319-
symbols: symbols_strings,
320-
level,
321-
});
322-
names.iter().rev().for_each(|name| self.store_name(&name));
321+
self.store_name(&name.symbol);
323322
}
324323
}
324+
325+
// Pop module from stack:
326+
self.emit(Instruction::Pop);
325327
}
326328
}
327329
ast::Statement::Expression { expression } => {

‎compiler/src/symboltable.rs‎

Lines changed: 5 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -296,24 +296,14 @@ impl SymbolTableBuilder {
296296
ast::Statement::Break | ast::Statement::Continue | ast::Statement::Pass => {
297297
// No symbols here.
298298
}
299-
ast::Statement::Import { import_parts } => {
300-
for part in import_parts {
301-
if let Some(alias) = &part.alias {
299+
ast::Statement::Import { names } | ast::Statement::ImportFrom { names, .. } => {
300+
for name in names {
301+
if let Some(alias) = &name.alias {
302302
// `import mymodule as myalias`
303303
self.register_name(alias, SymbolRole::Assigned)?;
304-
} else if part.symbols.is_empty() {
305-
// `import module`
306-
self.register_name(&part.module, SymbolRole::Assigned)?;
307304
} else {
308-
// `from mymodule import myimport`
309-
for symbol in &part.symbols {
310-
if let Some(alias) = &symbol.alias {
311-
// `from mymodule import myimportname as myalias`
312-
self.register_name(alias, SymbolRole::Assigned)?;
313-
} else {
314-
self.register_name(&symbol.symbol, SymbolRole::Assigned)?;
315-
}
316-
}
305+
// `import module`
306+
self.register_name(&name.symbol, SymbolRole::Assigned)?;
317307
}
318308
}
319309
}

‎parser/src/ast.rs‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,6 @@ pub struct ImportSymbol {
3232
pub alias: Option<String>,
3333
}
3434

35-
#[derive(Debug, PartialEq)]
36-
pub struct SingleImport {
37-
pub module: String,
38-
pub alias: Option<String>,
39-
pub symbols: Vec<ImportSymbol>,
40-
pub level: usize,
41-
}
42-
4335
#[derive(Debug, PartialEq)]
4436
pub struct Located<T> {
4537
pub location: Location,
@@ -58,7 +50,12 @@ pub enum Statement {
5850
value: Option<Expression>,
5951
},
6052
Import {
61-
import_parts: Vec<SingleImport>,
53+
names: Vec<ImportSymbol>,
54+
},
55+
ImportFrom {
56+
level: usize,
57+
module: Option<String>,
58+
names: Vec<ImportSymbol>,
6259
},
6360
Pass,
6461
Assert {

‎parser/src/python.lalrpop‎

Lines changed: 23 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -199,67 +199,46 @@ RaiseStatement: ast::LocatedStatement = {
199199
};
200200

201201
ImportStatement: ast::LocatedStatement = {
202-
<loc:@L> "import" <i: Comma<ImportPart<<DottedName>>>> => {
203-
ast::LocatedStatement {
204-
location: loc,
205-
node: ast::Statement::Import {
206-
import_parts: i
207-
.iter()
208-
.map(|(n, a)|
209-
ast::SingleImport {
210-
module: n.to_string(),
211-
symbols: vec![],
212-
alias: a.clone(),
213-
level: 0,
214-
})
215-
.collect()
216-
},
217-
}
218-
},
219-
<loc:@L> "from" <n:ImportFromLocation> "import" <i: ImportAsNames> => {
220-
ast::LocatedStatement {
221-
location: loc,
222-
node: ast::Statement::Import {
223-
import_parts: vec![
224-
ast::SingleImport {
225-
module: n.0.to_string(),
226-
symbols: i.iter()
227-
.map(|(i, a)|
228-
ast::ImportSymbol {
229-
symbol: i.to_string(),
230-
alias: a.clone(),
231-
})
232-
.collect(),
233-
alias: None,
234-
level: n.1
235-
}]
236-
},
237-
}
238-
},
202+
<loc:@L> "import" <names: Comma<ImportPart<<DottedName>>>> => {
203+
ast::LocatedStatement {
204+
location: loc,
205+
node: ast::Statement::Import { names },
206+
}
207+
},
208+
<loc:@L> "from" <n:ImportFromLocation> "import" <names: ImportAsNames> => {
209+
ast::LocatedStatement {
210+
location: loc,
211+
node: ast::Statement::ImportFrom {
212+
level: n.0,
213+
module: n.1,
214+
names
215+
},
216+
}
217+
},
239218
};
240219

241-
ImportFromLocation: (String, usize) = {
220+
ImportFromLocation: (usize, Option<String>) = {
242221
<dots: "."*> <name:DottedName> => {
243-
(name, dots.len())
222+
(dots.len(), Some(name))
244223
},
245224
<dots: "."+> => {
246-
("".to_string(), dots.len())
225+
(dots.len(), None)
247226
},
248227
};
249228

250-
ImportAsNames: Vec<(String, Option<String>)> = {
229+
ImportAsNames: Vec<ast::ImportSymbol> = {
251230
<i:Comma<ImportPart<Identifier>>> => i,
252231
"(" <i:Comma<ImportPart<Identifier>>> ")" => i,
253232
"*" => {
254233
// Star import all
255-
vec![("*".to_string(), None)]
234+
vec![ast::ImportSymbol { symbol: "*".to_string(), alias: None }]
256235
},
257236
};
258237

259238

260239
#[inline]
261-
ImportPart<I>: (String, Option<String>) = {
262-
<i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
240+
ImportPart<I>: ast::ImportSymbol = {
241+
<symbol:I> <a: ("as" Identifier)?> => ast::ImportSymbol { symbol, alias: a.map(|a| a.1) },
263242
};
264243

265244
// A name like abc or abc.def.ghi

‎vm/src/frame.rs‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,7 @@ impl Frame {
370370
ref name,
371371
ref level,
372372
} => self.import_star(vm, name, *level),
373+
bytecode::Instruction::ImportFrom { ref name } => self.import_from(vm, name),
373374
bytecode::Instruction::LoadName {
374375
ref name,
375376
ref scope,
@@ -932,16 +933,18 @@ impl Frame {
932933
.collect();
933934
let module = vm.import(module, &vm.ctx.new_tuple(from_list), level)?;
934935

935-
if symbols.is_empty() {
936-
self.push_value(module);
937-
} else {
938-
for symbol in symbols {
939-
let obj = vm
940-
.get_attribute(module.clone(), symbol.as_str())
941-
.map_err(|_| vm.new_import_error(format!("cannot import name '{}'", symbol)));
942-
self.push_value(obj?);
943-
}
944-
}
936+
self.push_value(module);
937+
Ok(None)
938+
}
939+
940+
#[cfg_attr(feature = "flame-it", flame("Frame"))]
941+
fn import_from(&self, vm: &VirtualMachine, name: &str) -> FrameResult {
942+
let module = self.last_value();
943+
// Load attribute, and transform any error into import error.
944+
let obj = vm
945+
.get_attribute(module, name)
946+
.map_err(|_| vm.new_import_error(format!("cannot import name '{}'", name)));
947+
self.push_value(obj?);
945948
Ok(None)
946949
}
947950

‎vm/src/stdlib/ast.rs‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,18 @@ fn statement_to_ast(
180180
ast::Statement::Expression { expression } => node!(vm, Expr, {
181181
value => expression_to_ast(vm, expression)?
182182
}),
183-
ast::Statement::Import { .. } => node!(vm, Import),
183+
ast::Statement::Import { names } => node!(vm, Import, {
184+
names => map_ast(alias_to_ast, vm, names)?
185+
}),
186+
ast::Statement::ImportFrom {
187+
level,
188+
module,
189+
names,
190+
} => node!(vm, ImportFrom, {
191+
level => vm.ctx.new_int(*level),
192+
module => optional_string_to_py_obj(vm, module),
193+
names => map_ast(alias_to_ast, vm, names)?
194+
}),
184195
ast::Statement::Nonlocal { names } => node!(vm, Nonlocal, {
185196
names => make_string_list(vm, names)
186197
}),
@@ -209,6 +220,13 @@ fn statement_to_ast(
209220
Ok(node)
210221
}
211222

223+
fn alias_to_ast(vm: &VirtualMachine, alias: &ast::ImportSymbol) -> PyResult<AstNodeRef> {
224+
Ok(node!(vm, alias, {
225+
symbol => vm.ctx.new_str(alias.symbol.to_string()),
226+
alias => optional_string_to_py_obj(vm, &alias.alias)
227+
}))
228+
}
229+
212230
fn optional_statements_to_ast(
213231
vm: &VirtualMachine,
214232
statements: &Option<Vec<ast::LocatedStatement>>,
@@ -610,6 +628,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
610628
"If" => py_class!(ctx, "_ast.If", ast_base.clone(), {}),
611629
"IfExp" => py_class!(ctx, "_ast.IfExp", ast_base.clone(), {}),
612630
"Import" => py_class!(ctx, "_ast.Import", ast_base.clone(), {}),
631+
"ImportFrom" => py_class!(ctx, "_ast.ImportFrom", ast_base.clone(), {}),
613632
"JoinedStr" => py_class!(ctx, "_ast.JoinedStr", ast_base.clone(), {}),
614633
"keyword" => py_class!(ctx, "_ast.keyword", ast_base.clone(), {}),
615634
"Lambda" => py_class!(ctx, "_ast.Lambda", ast_base.clone(), {}),

0 commit comments

Comments
 (0)