Calculate import level at parsing

compiler/src/bytecode.rs

@@ -51,9 +51,11 @@ pub enum Instruction {
     Import {
         name: String,
         symbols: Vec<String>,
+        level: usize,
     },
     ImportStar {
         name: String,
+        level: usize,
     },
     LoadName {
         name: String,
@@ -327,11 +329,25 @@ impl Instruction {
             ($variant:ident, $var1:expr, $var2:expr) => {
                 write!(f, "{:20} ({}, {})\n", stringify!($variant), $var1, $var2)
             };
+            ($variant:ident, $var1:expr, $var2:expr, $var3:expr) => {
+                write!(
+                    f,
+                    "{:20} ({}, {}, {})\n",
+                    stringify!($variant),
+                    $var1,
+                    $var2,
+                    $var3
+                )
+            };
         }
 
         match self {
-            Import { name, symbols } => w!(Import, name, format!("{:?}", symbols)),
-            ImportStar { name } => w!(ImportStar, name),
+            Import {
+                name,
+                symbols,
+                level,
+            } => w!(Import, name, format!("{:?}", symbols), level),
+            ImportStar { name, level } => w!(ImportStar, name, level),
             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

@@ -252,13 +252,16 @@ impl Compiler {
                     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() {
@@ -266,6 +269,7 @@ impl Compiler {
                         self.emit(Instruction::Import {
                             name: module.clone(),
                             symbols: vec![],
+                            level,
                         });
                         self.store_name(&module.clone());
                     } else {
@@ -276,6 +280,7 @@ impl Compiler {
                             // from module import *
                             self.emit(Instruction::ImportStar {
                                 name: module.clone(),
+                                level,
                             });
                         } else {
                             // from module import symbol
@@ -292,6 +297,7 @@ impl Compiler {
                             self.emit(Instruction::Import {
                                 name: module.clone(),
                                 symbols: symbols_strings,
+                                level,
                             });
                             names.iter().rev().for_each(|name| self.store_name(&name));
                         }

parser/src/ast.rs

@@ -38,6 +38,7 @@ pub struct SingleImport {
     pub module: String,
     pub alias: Option<String>,
     pub symbols: Vec<ImportSymbol>,
+    pub level: usize,
 }
 
 #[derive(Debug, PartialEq)]

parser/src/python.lalrpop

@@ -207,7 +207,8 @@ ImportStatement: ast::LocatedStatement = {
                 ast::SingleImport {
                     module: n.to_string(),
                     symbols: vec![],
-                    alias: a.clone()
+                    alias: a.clone(),
+                    level: 0,
                 })
             .collect()
       },
@@ -219,7 +220,7 @@ ImportStatement: ast::LocatedStatement = {
       node: ast::Statement::Import {
         import_parts: vec![
             ast::SingleImport {
-                module: n.to_string(),
+                module: n.0.to_string(),
                 symbols: i.iter()
                     .map(|(i, a)|
                         ast::ImportSymbol {
@@ -227,28 +228,20 @@ ImportStatement: ast::LocatedStatement = {
                             alias: a.clone(),
                         })
                     .collect(),
-                alias: None
+                alias: None,
+                level: n.1
         }]
       },
     }
   },
 };
 
-ImportFromLocation: String = {
+ImportFromLocation: (String, usize) = {
     <dots: "."*> <name:DottedName> => {
-       let mut r = "".to_string();
-       for _dot in dots {
-         r.push_str(".");
-       }
-      r.push_str(&name);
-      r
+       (name, dots.len())
     },
     <dots: "."+> => {
-       let mut r = "".to_string();
-       for _dot in dots {
-         r.push_str(".");
-       }
-      r
+       ("".to_string(), dots.len())
     },
 };
 

vm/src/frame.rs

@@ -358,8 +358,12 @@ impl Frame {
             bytecode::Instruction::Import {
                 ref name,
                 ref symbols,
-            } => self.import(vm, name, symbols),
-            bytecode::Instruction::ImportStar { ref name } => self.import_star(vm, name),
+                ref level,
+            } => self.import(vm, name, symbols, level),
+            bytecode::Instruction::ImportStar {
+                ref name,
+                ref level,
+            } => self.import_star(vm, name, level),
             bytecode::Instruction::LoadName {
                 ref name,
                 ref scope,
@@ -907,14 +911,18 @@ impl Frame {
         }
     }
 
-    fn import(&self, vm: &VirtualMachine, module: &str, symbols: &Vec<String>) -> FrameResult {
+    fn import(
+        &self,
+        vm: &VirtualMachine,
+        module: &str,
+        symbols: &Vec<String>,
+        level: &usize,
+    ) -> FrameResult {
         let from_list = symbols
             .iter()
             .map(|symbol| vm.ctx.new_str(symbol.to_string()))
             .collect();
-        let level = module.chars().take_while(|char| *char == '.').count();
-        let module_name = &module[level..];
-        let module = vm.import(module_name, &vm.ctx.new_tuple(from_list), level)?;
+        let module = vm.import(module, &vm.ctx.new_tuple(from_list), *level)?;
 
         if symbols.is_empty() {
             self.push_value(module);
@@ -929,9 +937,8 @@ impl Frame {
         Ok(None)
     }
 
-    fn import_star(&self, vm: &VirtualMachine, module: &str) -> FrameResult {
-        let level = module.chars().take_while(|char| *char == '.').count();
-        let module = vm.import(module, &vm.ctx.new_tuple(vec![]), level)?;
+    fn import_star(&self, vm: &VirtualMachine, module: &str, level: &usize) -> FrameResult {
+        let module = vm.import(module, &vm.ctx.new_tuple(vec![]), *level)?;
 
         // Grab all the names from the module and put them in the context
         if let Some(dict) = &module.dict {