Skip to content

Commit 816038c

Browse files
Merge pull request RustPython#54 from OddBloke/import
Implement complex import handling
2 parents a7f1830 + a7428e3 commit 816038c

9 files changed

Lines changed: 131 additions & 22 deletions

File tree

parser/src/ast.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,14 @@ pub struct Program {
2020
pub statements: Vec<Statement>,
2121
}
2222

23+
#[derive(Debug, PartialEq)]
24+
pub struct SingleImport {
25+
pub module: String,
26+
// (symbol name in module, name it should be assigned locally)
27+
pub symbol: Option<String>,
28+
pub alias: Option<String>,
29+
}
30+
2331
#[derive(Debug, PartialEq)]
2432
pub enum Statement {
2533
Break,
@@ -28,7 +36,7 @@ pub enum Statement {
2836
value: Option<Vec<Expression>>,
2937
},
3038
Import {
31-
name: String,
39+
import_parts: Vec<SingleImport>,
3240
},
3341
Pass,
3442
Assert {

parser/src/python.lalrpop

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,37 @@ FlowStatement: ast::Statement = {
9898
};
9999

100100
ImportStatement: ast::Statement = {
101-
"import" <n:DottedName> => ast::Statement::Import { name: n },
101+
"import" <i: Comma<ImportPart<<DottedName>>>> => {
102+
ast::Statement::Import {
103+
import_parts: i
104+
.iter()
105+
.map(|(n, a)|
106+
ast::SingleImport {
107+
module: n.to_string(),
108+
symbol: None,
109+
alias: a.clone()
110+
})
111+
.collect()
112+
}
113+
},
114+
"from" <n:DottedName> "import" <i: Comma<ImportPart<Identifier>>> => {
115+
ast::Statement::Import {
116+
import_parts: i
117+
.iter()
118+
.map(|(i, a)|
119+
ast::SingleImport {
120+
module: n.to_string(),
121+
symbol: Some(i.to_string()),
122+
alias: a.clone()
123+
})
124+
.collect()
125+
}
126+
},
127+
};
128+
129+
#[inline]
130+
ImportPart<I>: (String, Option<String>) = {
131+
<i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
102132
};
103133

104134
DottedName: String = {
@@ -436,6 +466,7 @@ extern {
436466
"in" => lexer::Tok::In,
437467
"is" => lexer::Tok::Is,
438468
"import" => lexer::Tok::Import,
469+
"from" => lexer::Tok::From,
439470
"not" => lexer::Tok::Not,
440471
"or" => lexer::Tok::Or,
441472
"pass" => lexer::Tok::Pass,

tests/snippets/import.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import import_target, import_target as aliased
2+
from import_target import func, other_func
3+
from import_target import func as aliased_func, other_func as aliased_other_func
4+
5+
assert import_target.X == import_target.func()
6+
assert import_target.X == func()
7+
8+
assert import_target.Y == other_func()
9+
10+
assert import_target.X == aliased.X
11+
assert import_target.Y == aliased.Y
12+
13+
assert import_target.X == aliased_func()
14+
assert import_target.Y == aliased_other_func()
15+
16+
#try:
17+
# X
18+
#except NameError:
19+
# pass
20+
#else:
21+
# raise AssertionError('X should not be imported')

tests/snippets/import_target.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# This is used by import.py; the two should be modified in concert
2+
3+
X = '123'
4+
Y = 'abc'
5+
6+
def func():
7+
return X
8+
9+
def other_func():
10+
return Y

tests/test_snippets.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,9 @@ def perform_test(filename, method, test_type):
5353

5454
def run_via_cpython(filename):
5555
""" Simply invoke python itself on the script """
56-
subprocess.check_call([sys.executable, filename])
56+
env = os.environ.copy()
57+
env['PYTHONPATH'] = '.'
58+
subprocess.check_call([sys.executable, filename], env=env)
5759

5860

5961
def run_via_cpython_bytecode(filename, test_type):
@@ -76,8 +78,12 @@ def run_via_rustpython(filename, test_type):
7678
log_level = 'info' if test_type == _TestType.benchmark else 'trace'
7779
env['RUST_LOG'] = '{},cargo=error,jobserver=error'.format(log_level)
7880
env['RUST_BACKTRACE'] = '1'
81+
# XXX: Once we support PYTHONPATH (or similar), we should use that instead
82+
# of changing directory
83+
cwd = os.path.dirname(filename)
7984
with pushd(RUSTPYTHON_RUNNER_DIR):
80-
subprocess.check_call(['cargo', 'run', '--release', filename], env=env)
85+
subprocess.check_call(
86+
['cargo', 'run', '--release', filename], env=env, cwd=cwd)
8187

8288

8389
def create_test_function(cls, filename, method, test_type):

vm/src/bytecode.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,10 @@ pub type Label = usize;
3434

3535
#[derive(Debug, Clone)]
3636
pub enum Instruction {
37-
Import { name: String },
37+
Import {
38+
name: String,
39+
symbol: Option<String>,
40+
},
3841
LoadName { name: String },
3942
StoreName { name: String },
4043
StoreSubscript,

vm/src/compile.rs

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,27 @@ impl Compiler {
106106
fn compile_statement(&mut self, statement: &ast::Statement) {
107107
trace!("Compiling {:?}", statement);
108108
match statement {
109-
ast::Statement::Import { name } => {
110-
self.emit(Instruction::Import { name: name.clone() });
111-
self.emit(Instruction::StoreName { name: name.clone() });
109+
ast::Statement::Import { import_parts } => {
110+
for ast::SingleImport {
111+
module,
112+
symbol,
113+
alias,
114+
} in import_parts
115+
{
116+
self.emit(Instruction::Import {
117+
name: module.clone(),
118+
symbol: symbol.clone().map(|s| s.clone()),
119+
});
120+
self.emit(Instruction::StoreName {
121+
name: match alias {
122+
Some(alias) => alias.clone(),
123+
None => match symbol {
124+
Some(symbol) => symbol.clone(),
125+
None => module.clone(),
126+
},
127+
},
128+
});
129+
}
112130
}
113131
ast::Statement::Expression { expression } => {
114132
self.compile_expression(expression);

vm/src/import.rs

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,15 @@ use super::compile;
1313
use super::pyobject::{PyObject, PyObjectKind, PyResult, DictProtocol};
1414
use super::vm::VirtualMachine;
1515

16-
pub fn import(vm: &mut VirtualMachine, name: &String) -> PyResult {
16+
fn import_module(vm: &mut VirtualMachine, module: &String) -> PyResult {
1717
// First, see if we already loaded the module:
1818
let sys_modules = vm.sys_module.get_item(&"modules".to_string());
19-
if sys_modules.contains_key(name) {
20-
return Ok(sys_modules.get_item(name))
19+
if sys_modules.contains_key(module) {
20+
return Ok(sys_modules.get_item(module));
2121
}
2222

2323
// Time to search for module in any place:
24-
let filepath = find_source(name).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
24+
let filepath = find_source(module).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
2525
let source = parser::read_file(filepath.as_path())
2626
.map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
2727

@@ -42,14 +42,23 @@ pub fn import(vm: &mut VirtualMachine, name: &String) -> PyResult {
4242
Ok(_) => {}
4343
Err(value) => return Err(value),
4444
}
45+
Ok(scope)
46+
}
4547

46-
let obj = PyObject::new(
47-
PyObjectKind::Module {
48-
name: name.clone(),
49-
dict: scope.clone(),
50-
},
51-
vm.get_type(),
52-
);
48+
pub fn import(vm: &mut VirtualMachine, module: &String, symbol: &Option<String>) -> PyResult {
49+
let scope = import_module(vm, module)?;
50+
// If we're importing a symbol, look it up and use it, otherwise construct a module and return
51+
// that
52+
let obj = match symbol {
53+
Some(symbol) => scope.get_item(symbol),
54+
None => PyObject::new(
55+
PyObjectKind::Module {
56+
name: module.clone(),
57+
dict: scope.clone(),
58+
},
59+
vm.get_type(),
60+
),
61+
};
5362
Ok(obj)
5463
}
5564

vm/src/vm.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -510,8 +510,8 @@ impl VirtualMachine {
510510
}
511511
}
512512

513-
fn import(&mut self, name: &String) -> Option<PyResult> {
514-
let obj = match import(self, name) {
513+
fn import(&mut self, module: &String, symbol: &Option<String>) -> Option<PyResult> {
514+
let obj = match import(self, module, symbol) {
515515
Ok(value) => value,
516516
Err(value) => return Some(Err(value)),
517517
};
@@ -564,7 +564,10 @@ impl VirtualMachine {
564564
self.push_value(obj);
565565
None
566566
}
567-
bytecode::Instruction::Import { ref name } => self.import(name),
567+
bytecode::Instruction::Import {
568+
ref name,
569+
ref symbol,
570+
} => self.import(name, symbol),
568571
bytecode::Instruction::LoadName { ref name } => self.load_name(name),
569572
bytecode::Instruction::StoreName { ref name } => {
570573
// take top of stack and assign in scope:

0 commit comments

Comments
 (0)