Merge pull request #1151 from RustPython/import-syntax

Simplify import AST in line with CPython.

bytecode/src/bytecode.rs

@@ -88,6 +88,9 @@ pub enum Instruction {
         name: String,
         level: usize,
     },
+    ImportFrom {
+        name: String,
+    },
     LoadName {
         name: String,
         scope: NameScope,
@@ -379,6 +382,7 @@ impl Instruction {
                 level,
             } => w!(Import, name, format!("{:?}", symbols), level),
             ImportStar { name, level } => w!(ImportStar, name, level),
+            ImportFrom { name } => w!(ImportFrom, name),
             LoadName { name, scope } => w!(LoadName, name, format!("{:?}", scope)),
             StoreName { name, scope } => w!(StoreName, name, format!("{:?}", scope)),
             DeleteName { name } => w!(DeleteName, name),

compiler/src/compile.rs

@@ -267,61 +267,63 @@ impl Compiler {
         self.set_source_location(&statement.location);
 
         match &statement.node {
-            ast::Statement::Import { import_parts } => {
-                for ast::SingleImport {
-                    module,
-                    symbols,
-                    alias,
-                    level,
-                } in import_parts
-                {
-                    let level = *level;
-                    if let Some(alias) = alias {
-                        // import module as alias
-                        self.emit(Instruction::Import {
-                            name: module.clone(),
-                            symbols: vec![],
-                            level,
-                        });
-                        self.store_name(&alias);
-                    } else if symbols.is_empty() {
-                        // import module
-                        self.emit(Instruction::Import {
-                            name: module.clone(),
-                            symbols: vec![],
-                            level,
-                        });
-                        self.store_name(&module.clone());
+            ast::Statement::Import { names } => {
+                // import a, b, c as d
+                for name in names {
+                    self.emit(Instruction::Import {
+                        name: name.symbol.clone(),
+                        symbols: vec![],
+                        level: 0,
+                    });
+
+                    if let Some(alias) = &name.alias {
+                        self.store_name(alias);
                     } else {
-                        let import_star = symbols
-                            .iter()
-                            .any(|import_symbol| import_symbol.symbol == "*");
-                        if import_star {
-                            // from module import *
-                            self.emit(Instruction::ImportStar {
-                                name: module.clone(),
-                                level,
-                            });
+                        self.store_name(&name.symbol);
+                    }
+                }
+            }
+            ast::Statement::ImportFrom {
+                level,
+                module,
+                names,
+            } => {
+                let import_star = names.iter().any(|n| n.symbol == "*");
+
+                if import_star {
+                    // from .... import *
+                    self.emit(Instruction::ImportStar {
+                        name: module.clone().unwrap(),
+                        level: *level,
+                    });
+                } else {
+                    // from mod import a, b as c
+                    // First, determine the fromlist (for import lib):
+                    let from_list = names.iter().map(|n| n.symbol.clone()).collect();
+
+                    // Load module once:
+                    self.emit(Instruction::Import {
+                        name: module.clone().unwrap(),
+                        symbols: from_list,
+                        level: *level,
+                    });
+
+                    for name in names {
+                        // import symbol from module:
+                        self.emit(Instruction::ImportFrom {
+                            name: name.symbol.to_string(),
+                        });
+
+                        // Store module under proper name:
+                        if let Some(alias) = &name.alias {
+                            self.store_name(alias);
                         } else {
-                            // from module import symbol
-                            // from module import symbol as alias
-                            let (names, symbols_strings): (Vec<String>, Vec<String>) = symbols
-                                .iter()
-                                .map(|ast::ImportSymbol { symbol, alias }| {
-                                    (
-                                        alias.clone().unwrap_or_else(|| symbol.to_string()),
-                                        symbol.to_string(),
-                                    )
-                                })
-                                .unzip();
-                            self.emit(Instruction::Import {
-                                name: module.clone(),
-                                symbols: symbols_strings,
-                                level,
-                            });
-                            names.iter().rev().for_each(|name| self.store_name(&name));
+                            self.store_name(&name.symbol);
                         }
                     }
+
+                    // Pop module from stack:
+                    self.emit(Instruction::Pop);
                 }
             }
             ast::Statement::Expression { expression } => {

compiler/src/symboltable.rs

@@ -296,24 +296,14 @@ impl SymbolTableBuilder {
             ast::Statement::Break | ast::Statement::Continue | ast::Statement::Pass => {
                 // No symbols here.
             }
-            ast::Statement::Import { import_parts } => {
-                for part in import_parts {
-                    if let Some(alias) = &part.alias {
+            ast::Statement::Import { names } | ast::Statement::ImportFrom { names, .. } => {
+                for name in names {
+                    if let Some(alias) = &name.alias {
                         // `import mymodule as myalias`
                         self.register_name(alias, SymbolRole::Assigned)?;
-                    } else if part.symbols.is_empty() {
-                        // `import module`
-                        self.register_name(&part.module, SymbolRole::Assigned)?;
                     } else {
-                        // `from mymodule import myimport`
-                        for symbol in &part.symbols {
-                            if let Some(alias) = &symbol.alias {
-                                // `from mymodule import myimportname as myalias`
-                                self.register_name(alias, SymbolRole::Assigned)?;
-                            } else {
-                                self.register_name(&symbol.symbol, SymbolRole::Assigned)?;
-                            }
-                        }
+                        // `import module`
+                        self.register_name(&name.symbol, SymbolRole::Assigned)?;
                     }
                 }
             }

parser/src/ast.rs

@@ -32,14 +32,6 @@ pub struct ImportSymbol {
     pub alias: Option<String>,
 }
 
-#[derive(Debug, PartialEq)]
-pub struct SingleImport {
-    pub module: String,
-    pub alias: Option<String>,
-    pub symbols: Vec<ImportSymbol>,
-    pub level: usize,
-}
-
 #[derive(Debug, PartialEq)]
 pub struct Located<T> {
     pub location: Location,
@@ -58,7 +50,12 @@ pub enum Statement {
         value: Option<Expression>,
     },
     Import {
-        import_parts: Vec<SingleImport>,
+        names: Vec<ImportSymbol>,
+    },
+    ImportFrom {
+        level: usize,
+        module: Option<String>,
+        names: Vec<ImportSymbol>,
     },
     Pass,
     Assert {

parser/src/python.lalrpop

@@ -199,67 +199,46 @@ RaiseStatement: ast::LocatedStatement = {
 };
 
 ImportStatement: ast::LocatedStatement = {
-  <loc:@L> "import" <i: Comma<ImportPart<<DottedName>>>> => {
-    ast::LocatedStatement {
-      location: loc,
-      node: ast::Statement::Import {
-        import_parts: i
-            .iter()
-            .map(|(n, a)|
-                ast::SingleImport {
-                    module: n.to_string(),
-                    symbols: vec![],
-                    alias: a.clone(),
-                    level: 0,
-                })
-            .collect()
-      },
-    }
-  },
-  <loc:@L> "from" <n:ImportFromLocation> "import" <i: ImportAsNames> => {
-    ast::LocatedStatement {
-      location: loc,
-      node: ast::Statement::Import {
-        import_parts: vec![
-            ast::SingleImport {
-                module: n.0.to_string(),
-                symbols: i.iter()
-                    .map(|(i, a)|
-                        ast::ImportSymbol {
-                            symbol: i.to_string(),
-                            alias: a.clone(),
-                        })
-                    .collect(),
-                alias: None,
-                level: n.1
-        }]
-      },
-    }
-  },
+    <loc:@L> "import" <names: Comma<ImportPart<<DottedName>>>> => {
+        ast::LocatedStatement {
+            location: loc,
+            node: ast::Statement::Import { names },
+        }
+    },
+    <loc:@L> "from" <n:ImportFromLocation> "import" <names: ImportAsNames> => {
+        ast::LocatedStatement {
+            location: loc,
+            node: ast::Statement::ImportFrom {
+                level: n.0,
+                module: n.1,
+                names
+            },
+        }
+    },
 };
 
-ImportFromLocation: (String, usize) = {
+ImportFromLocation: (usize, Option<String>) = {
     <dots: "."*> <name:DottedName> => {
-       (name, dots.len())
+        (dots.len(), Some(name))
     },
     <dots: "."+> => {
-       ("".to_string(), dots.len())
+        (dots.len(), None)
     },
 };
 
-ImportAsNames: Vec<(String, Option<String>)> = {
+ImportAsNames: Vec<ast::ImportSymbol> = {
     <i:Comma<ImportPart<Identifier>>> => i,
     "(" <i:Comma<ImportPart<Identifier>>> ")" => i,
     "*" => {
         // Star import all
-        vec![("*".to_string(), None)]
+        vec![ast::ImportSymbol { symbol: "*".to_string(), alias: None }]
     },
 };
 
 
 #[inline]
-ImportPart<I>: (String, Option<String>) = {
-    <i:I> <a: ("as" Identifier)?> => (i, a.map(|a| a.1)),
+ImportPart<I>: ast::ImportSymbol = {
+    <symbol:I> <a: ("as" Identifier)?> => ast::ImportSymbol { symbol, alias: a.map(|a| a.1) },
 };
 
 // A name like abc or abc.def.ghi

vm/src/frame.rs

@@ -370,6 +370,7 @@ impl Frame {
                 ref name,
                 ref level,
             } => self.import_star(vm, name, *level),
+            bytecode::Instruction::ImportFrom { ref name } => self.import_from(vm, name),
             bytecode::Instruction::LoadName {
                 ref name,
                 ref scope,
@@ -932,16 +933,18 @@ impl Frame {
             .collect();
         let module = vm.import(module, &vm.ctx.new_tuple(from_list), level)?;
 
-        if symbols.is_empty() {
-            self.push_value(module);
-        } else {
-            for symbol in symbols {
-                let obj = vm
-                    .get_attribute(module.clone(), symbol.as_str())
-                    .map_err(|_| vm.new_import_error(format!("cannot import name '{}'", symbol)));
-                self.push_value(obj?);
-            }
-        }
+        self.push_value(module);
+        Ok(None)
+    }
+
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
+    fn import_from(&self, vm: &VirtualMachine, name: &str) -> FrameResult {
+        let module = self.last_value();
+        // Load attribute, and transform any error into import error.
+        let obj = vm
+            .get_attribute(module, name)
+            .map_err(|_| vm.new_import_error(format!("cannot import name '{}'", name)));
+        self.push_value(obj?);
         Ok(None)
     }
 

vm/src/stdlib/ast.rs

@@ -180,7 +180,18 @@ fn statement_to_ast(
         ast::Statement::Expression { expression } => node!(vm, Expr, {
             value => expression_to_ast(vm, expression)?
         }),
-        ast::Statement::Import { .. } => node!(vm, Import),
+        ast::Statement::Import { names } => node!(vm, Import, {
+            names => map_ast(alias_to_ast, vm, names)?
+        }),
+        ast::Statement::ImportFrom {
+            level,
+            module,
+            names,
+        } => node!(vm, ImportFrom, {
+            level => vm.ctx.new_int(*level),
+            module => optional_string_to_py_obj(vm, module),
+            names => map_ast(alias_to_ast, vm, names)?
+        }),
         ast::Statement::Nonlocal { names } => node!(vm, Nonlocal, {
             names => make_string_list(vm, names)
         }),
@@ -209,6 +220,13 @@ fn statement_to_ast(
     Ok(node)
 }
 
+fn alias_to_ast(vm: &VirtualMachine, alias: &ast::ImportSymbol) -> PyResult<AstNodeRef> {
+    Ok(node!(vm, alias, {
+        symbol => vm.ctx.new_str(alias.symbol.to_string()),
+        alias => optional_string_to_py_obj(vm, &alias.alias)
+    }))
+}
+
 fn optional_statements_to_ast(
     vm: &VirtualMachine,
     statements: &Option<Vec<ast::LocatedStatement>>,
@@ -610,6 +628,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
         "If" => py_class!(ctx, "_ast.If", ast_base.clone(), {}),
         "IfExp" => py_class!(ctx, "_ast.IfExp", ast_base.clone(), {}),
         "Import" => py_class!(ctx, "_ast.Import", ast_base.clone(), {}),
+        "ImportFrom" => py_class!(ctx, "_ast.ImportFrom", ast_base.clone(), {}),
         "JoinedStr" => py_class!(ctx, "_ast.JoinedStr", ast_base.clone(), {}),
         "keyword" => py_class!(ctx, "_ast.keyword", ast_base.clone(), {}),
         "Lambda" => py_class!(ctx, "_ast.Lambda", ast_base.clone(), {}),