Merge pull request RustPython#839 from RustPython/nonlocal

Add nonlocal support.

Author: windelbouwman

tests/snippets/global_nonlocal.py

@@ -11,3 +11,13 @@ def b():
 b()
 assert a == 4
 
+def x():
+    def y():
+        nonlocal b
+        b = 3
+    b = 2
+    y()
+    return b
+
+res = x()
+assert res == 3, str(res)

vm/src/bytecode.rs

@@ -41,6 +41,7 @@ pub type Label = usize;
 #[derive(Debug, Clone, PartialEq)]
 pub enum NameScope {
     Local,
+    NonLocal,
     Global,
 }
 

vm/src/compile.rs

@@ -179,12 +179,12 @@ impl Compiler {
         let role = self.lookup_name(name);
         match role {
             SymbolRole::Global => bytecode::NameScope::Global,
+            SymbolRole::Nonlocal => bytecode::NameScope::NonLocal,
             _ => bytecode::NameScope::Local,
         }
     }
 
     fn load_name(&mut self, name: &str) {
-        // TODO: if global, do something else!
         let scope = self.scope_for_name(name);
         self.emit(Instruction::LoadName {
             name: name.to_string(),

vm/src/frame.rs

@@ -126,6 +126,7 @@ pub trait NameProtocol {
     fn store_name(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
     fn delete_name(&self, vm: &VirtualMachine, name: &str);
     fn load_cell(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
+    fn store_cell(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
     fn load_global(&self, vm: &VirtualMachine, name: &str) -> Option<PyObjectRef>;
     fn store_global(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef);
 }
@@ -154,6 +155,16 @@ impl NameProtocol for Scope {
         None
     }
 
+    fn store_cell(&self, vm: &VirtualMachine, name: &str, value: PyObjectRef) {
+        self.locals
+            .iter()
+            .skip(1)
+            .next()
+            .expect("no outer scope for non-local")
+            .set_item(name, value, vm)
+            .unwrap();
+    }
+
     fn store_name(&self, vm: &VirtualMachine, key: &str, value: PyObjectRef) {
         self.get_locals().set_item(key, value, vm).unwrap();
     }
@@ -1029,8 +1040,11 @@ impl Frame {
             bytecode::NameScope::Global => {
                 self.scope.store_global(vm, name, obj);
             }
+            bytecode::NameScope::NonLocal => {
+                self.scope.store_cell(vm, name, obj);
+            }
             bytecode::NameScope::Local => {
-                self.scope.store_name(&vm, name, obj);
+                self.scope.store_name(vm, name, obj);
             }
         }
         Ok(None)
@@ -1049,6 +1063,7 @@ impl Frame {
     ) -> FrameResult {
         let optional_value = match name_scope {
             bytecode::NameScope::Global => self.scope.load_global(vm, name),
+            bytecode::NameScope::NonLocal => self.scope.load_cell(vm, name),
             bytecode::NameScope::Local => self.scope.load_name(&vm, name),
         };