windelbouwman · Aug 10, 2018 · Aug 9, 2018 · Aug 9, 2018
diff --git a/parser/src/ast.rs b/parser/src/ast.rs
@@ -20,6 +20,14 @@ pub struct Program {
     pub statements: Vec<Statement>,
 }
 
+#[derive(Debug, PartialEq)]
+pub struct SingleImport {
+    pub module: String,
+    // (symbol name in module, name it should be assigned locally)
+    pub symbol: Option<String>,
+    pub alias: Option<String>,
+}
+
 #[derive(Debug, PartialEq)]
 pub enum Statement {
     Break,
@@ -28,7 +36,7 @@ pub enum Statement {
         value: Option<Vec<Expression>>,
     },
     Import {
-        name: String,
+        import_parts: Vec<SingleImport>,
     },
     Pass,
     Assert {

diff --git a/parser/src/python.lalrpop b/parser/src/python.lalrpop
@@ -98,7 +98,37 @@ FlowStatement: ast::Statement = {
 };
 
 ImportStatement: ast::Statement = {
-  "import" <n:DottedName> => ast::Statement::Import { name: n },
+  "import" <i: Comma<ImportPart<<DottedName>>>> => {
+    ast::Statement::Import {
+        import_parts: i
+            .iter()
+            .map(|(n, a)|
+                ast::SingleImport {
+                    module: n.to_string(),
+                    symbol: None,
+                    alias: a.clone()
+                })
+            .collect()
+    }
+  },
+  "from" <n:DottedName> "import" <i: Comma<ImportPart<Identifier>>> => {
+    ast::Statement::Import {
+        import_parts: i
+            .iter()
+            .map(|(i, a)|
+                ast::SingleImport {
+                    module: n.to_string(),
+                    symbol: Some(i.to_string()),
+                    alias: a.clone()
+                })
+            .collect()
+    }
+  },
+};
+
+#[inline]
+ImportPart<I>: (String, Option<String>) = {
+    <i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
 };
 
 DottedName: String = {
@@ -436,6 +466,7 @@ extern {
         "in" => lexer::Tok::In,
         "is" => lexer::Tok::Is,
         "import" => lexer::Tok::Import,
+        "from" => lexer::Tok::From,
         "not" => lexer::Tok::Not,
         "or" => lexer::Tok::Or,
         "pass" => lexer::Tok::Pass,

diff --git a/tests/snippets/import.py b/tests/snippets/import.py
@@ -0,0 +1,21 @@
+import import_target, import_target as aliased
+from import_target import func, other_func
+from import_target import func as aliased_func, other_func as aliased_other_func
+
+assert import_target.X == import_target.func()
+assert import_target.X == func()
+
+assert import_target.Y == other_func()
+
+assert import_target.X == aliased.X
+assert import_target.Y == aliased.Y
+
+assert import_target.X == aliased_func()
+assert import_target.Y == aliased_other_func()
+
+#try:
+#    X
+#except NameError:
+#    pass
+#else:
+#    raise AssertionError('X should not be imported')
diff --git a/tests/snippets/import_target.py b/tests/snippets/import_target.py
@@ -0,0 +1,10 @@
+# This is used by import.py; the two should be modified in concert
+
+X = '123'
+Y = 'abc'
+
+def func():
+    return X
+
+def other_func():
+    return Y
diff --git a/tests/test_snippets.py b/tests/test_snippets.py
@@ -53,7 +53,9 @@ def perform_test(filename, method, test_type):
 
 def run_via_cpython(filename):
     """ Simply invoke python itself on the script """
-    subprocess.check_call([sys.executable, filename])
+    env = os.environ.copy()
+    env['PYTHONPATH'] = '.'
+    subprocess.check_call([sys.executable, filename], env=env)
 
 
 def run_via_cpython_bytecode(filename, test_type):
@@ -76,8 +78,12 @@ def run_via_rustpython(filename, test_type):
     log_level = 'info' if test_type == _TestType.benchmark else 'trace'
     env['RUST_LOG'] = '{},cargo=error,jobserver=error'.format(log_level)
     env['RUST_BACKTRACE'] = '1'
+    # XXX: Once we support PYTHONPATH (or similar), we should use that instead
+    # of changing directory
+    cwd = os.path.dirname(filename)
     with pushd(RUSTPYTHON_RUNNER_DIR):
-        subprocess.check_call(['cargo', 'run', '--release', filename], env=env)
+        subprocess.check_call(
+            ['cargo', 'run', '--release', filename], env=env, cwd=cwd)
 
 
 def create_test_function(cls, filename, method, test_type):

diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs
@@ -34,7 +34,10 @@ pub type Label = usize;
 
 #[derive(Debug, Clone)]
 pub enum Instruction {
-    Import { name: String },
+    Import {
+        name: String,
+        symbol: Option<String>,
+    },
     LoadName { name: String },
     StoreName { name: String },
     StoreSubscript,

diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -106,9 +106,27 @@ impl Compiler {
     fn compile_statement(&mut self, statement: &ast::Statement) {
         trace!("Compiling {:?}", statement);
         match statement {
-            ast::Statement::Import { name } => {
-                self.emit(Instruction::Import { name: name.clone() });
-                self.emit(Instruction::StoreName { name: name.clone() });
+            ast::Statement::Import { import_parts } => {
+                for ast::SingleImport {
+                    module,
+                    symbol,
+                    alias,
+                } in import_parts
+                {
+                    self.emit(Instruction::Import {
+                        name: module.clone(),
+                        symbol: symbol.clone().map(|s| s.clone()),
+                    });
+                    self.emit(Instruction::StoreName {
+                        name: match alias {
+                            Some(alias) => alias.clone(),
+                            None => match symbol {
+                                Some(symbol) => symbol.clone(),
+                                None => module.clone(),
+                            },
+                        },
+                    });
+                }
             }
             ast::Statement::Expression { expression } => {
                 self.compile_expression(expression);

diff --git a/vm/src/import.rs b/vm/src/import.rs
@@ -13,15 +13,15 @@ use super::compile;
 use super::pyobject::{PyObject, PyObjectKind, PyResult, DictProtocol};
 use super::vm::VirtualMachine;
 
-pub fn import(vm: &mut VirtualMachine, name: &String) -> PyResult {
+fn import_module(vm: &mut VirtualMachine, module: &String) -> PyResult {
     // First, see if we already loaded the module:
     let sys_modules = vm.sys_module.get_item(&"modules".to_string());
-    if sys_modules.contains_key(name) {
-        return Ok(sys_modules.get_item(name))
+    if sys_modules.contains_key(module) {
+        return Ok(sys_modules.get_item(module));
     }
 
     // Time to search for module in any place:
-    let filepath = find_source(name).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
+    let filepath = find_source(module).map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
     let source = parser::read_file(filepath.as_path())
         .map_err(|e| vm.new_exception(format!("Error: {:?}", e)))?;
 
@@ -42,14 +42,23 @@ pub fn import(vm: &mut VirtualMachine, name: &String) -> PyResult {
         Ok(_) => {}
         Err(value) => return Err(value),
     }
+    Ok(scope)
+}
 
-    let obj = PyObject::new(
-        PyObjectKind::Module {
-            name: name.clone(),
-            dict: scope.clone(),
-        },
-        vm.get_type(),
-    );
+pub fn import(vm: &mut VirtualMachine, module: &String, symbol: &Option<String>) -> PyResult {
+    let scope = import_module(vm, module)?;
+    // If we're importing a symbol, look it up and use it, otherwise construct a module and return
+    // that
+    let obj = match symbol {
+        Some(symbol) => scope.get_item(symbol),
+        None => PyObject::new(
+            PyObjectKind::Module {
+                name: module.clone(),
+                dict: scope.clone(),
+            },
+            vm.get_type(),
+        ),
+    };
     Ok(obj)
 }
 

diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -510,8 +510,8 @@ impl VirtualMachine {
         }
     }
 
-    fn import(&mut self, name: &String) -> Option<PyResult> {
-        let obj = match import(self, name) {
+    fn import(&mut self, module: &String, symbol: &Option<String>) -> Option<PyResult> {
+        let obj = match import(self, module, symbol) {
             Ok(value) => value,
             Err(value) => return Some(Err(value)),
         };
@@ -564,7 +564,10 @@ impl VirtualMachine {
                 self.push_value(obj);
                 None
             }
-            bytecode::Instruction::Import { ref name } => self.import(name),
+            bytecode::Instruction::Import {
+                ref name,
+                ref symbol,
+            } => self.import(name, symbol),
             bytecode::Instruction::LoadName { ref name } => self.load_name(name),
             bytecode::Instruction::StoreName { ref name } => {
                 // take top of stack and assign in scope: