Merge pull request RustPython#237 from johndaniels/add-star-import

windelbouwman · web-flow · commit 946df5307713 · 2018-12-23T17:55:34.000+01:00
Add 'from x import *' syntax.
diff --git a/tests/snippets/import.py b/tests/snippets/import.py
@@ -1,6 +1,7 @@
 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
+from import_star import *
 
 assert import_target.X == import_target.func()
 assert import_target.X == func()
@@ -13,6 +14,8 @@
 assert import_target.X == aliased_func()
 assert import_target.Y == aliased_other_func()
 
+assert STAR_IMPORT == '123'
+
 # TODO: Once we can determine current directory, use that to construct this
 # path:
 #import sys
diff --git a/tests/snippets/import_star.py b/tests/snippets/import_star.py
@@ -0,0 +1,3 @@
+# This is used by import.py; the two should be modified in concert
+
+STAR_IMPORT = '123'
diff --git a/vm/src/bytecode.rs b/vm/src/bytecode.rs
@@ -66,6 +66,9 @@ pub enum Instruction {
         name: String,
         symbol: Option<String>,
     },
+    ImportStar {
+        name: String,
+    },
     LoadName {
         name: String,
     },
diff --git a/vm/src/compile.rs b/vm/src/compile.rs
@@ -156,19 +156,28 @@ impl Compiler {
                     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(),
-                            },
-                        },
-                    });
+                    match symbol {
+                        Some(name) if name == "*" => {
+                            self.emit(Instruction::ImportStar {
+                                name: module.clone(),
+                            });
+                        }
+                        _ => {
+                            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 } => {
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -7,7 +7,7 @@ use std::path::PathBuf;
 
 use super::builtins;
 use super::bytecode;
-use super::import::import;
+use super::import::{import, import_module};
 use super::obj::objbool;
 use super::obj::objcode;
 use super::obj::objdict;
@@ -182,6 +182,7 @@ impl Frame {
                 ref name,
                 ref symbol,
             } => self.import(vm, name, symbol),
+            bytecode::Instruction::ImportStar { ref name } => self.import_star(vm, name),
             bytecode::Instruction::LoadName { ref name } => self.load_name(vm, name),
             bytecode::Instruction::StoreName { ref name } => self.store_name(vm, name),
             bytecode::Instruction::DeleteName { ref name } => self.delete_name(vm, name),
@@ -677,6 +678,26 @@ impl Frame {
         Ok(None)
     }
 
+    fn import_star(&mut self, vm: &mut VirtualMachine, module: &str) -> FrameResult {
+        let current_path = match &self.code.source_path {
+            Some(source_path) => {
+                let mut source_pathbuf = PathBuf::from(source_path);
+                source_pathbuf.pop();
+                source_pathbuf
+            }
+            None => PathBuf::from("."),
+        };
+
+        // Grab all the names from the module and put them in the context
+        let obj = import_module(vm, current_path, module)?;
+
+        for (k, v) in obj.get_key_value_pairs().iter() {
+            vm.ctx
+                .set_item(&self.locals, &objstr::get_value(k), v.clone());
+        }
+        Ok(None)
+    }
+
     // Unwind all blocks:
     fn unwind_blocks(&mut self, vm: &mut VirtualMachine) -> Option<PyObjectRef> {
         loop {
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -25,7 +25,7 @@ use super::vm::VirtualMachine;
 use num_bigint::BigInt;
 use num_complex::Complex64;
 use num_traits::{One, Zero};
-use std::cell::RefCell;
+use std::cell::{Ref, RefCell};
 use std::collections::HashMap;
 use std::fmt;
 use std::rc::{Rc, Weak};
@@ -681,6 +681,7 @@ impl AttributeProtocol for PyObjectRef {
 pub trait DictProtocol {
     fn contains_key(&self, k: &str) -> bool;
     fn get_item(&self, k: &str) -> Option<PyObjectRef>;
+    fn get_key_value_pairs(&self) -> Vec<(PyObjectRef, PyObjectRef)>;
 }
 
 impl DictProtocol for PyObjectRef {
@@ -701,6 +702,15 @@ impl DictProtocol for PyObjectRef {
             _ => panic!("TODO"),
         }
     }
+
+    fn get_key_value_pairs(&self) -> Vec<(PyObjectRef, PyObjectRef)> {
+        match self.borrow().kind {
+            PyObjectKind::Dict { elements: _ } => objdict::get_key_value_pairs(self),
+            PyObjectKind::Module { name: _, ref dict } => dict.get_key_value_pairs(),
+            PyObjectKind::Scope { ref scope } => scope.locals.get_key_value_pairs(),
+            _ => panic!("TODO"),
+        }
+    }
 }
 
 impl fmt::Debug for PyObject {

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# This is used by import.py; the two should be modified in concert`
	`2`	`+`
	`3`	`+STAR_IMPORT = '123'`