py/compile: Fix scope of assignment expression target in comprehensions.

When := is used in a comprehension the target variable is bound to the
parent scope, so it's either a global or a nonlocal.  Prior to this commit
that was handled by simply using the parent scope's id_info for the
target variable.  That's completely wrong because it uses the slot number
for the parent's Python stack to store the variable, rather than the slot
number for the comprehension.  This will in most cases lead to incorrect
behaviour or memory faults.

This commit fixes the scoping of the target variable by explicitly
declaring it a global or nonlocal, depending on whether the parent is the
global scope or not.  Then the id_info of the comprehension can be used to
access the target variable.  This fixes a lot of cases of using := in a
comprehension.

Code size change for this commit:

       bare-arm:    +0 +0.000%
    minimal x86:    +0 +0.000%
       unix x64:  +152 +0.019% standard
          stm32:   +96 +0.024% PYBV10
         cc3200:   +96 +0.052%
        esp8266:  +196 +0.028% GENERIC
          esp32:  +156 +0.010% GENERIC[incl +8(data)]
         mimxrt:   +96 +0.027% TEENSY40
     renesas-ra:   +88 +0.014% RA6M2_EK
            nrf:   +88 +0.048% pca10040
            rp2:  +104 +0.020% PICO
           samd:   +88 +0.033% ADAFRUIT_ITSYBITSY_M4_EXPRESS

Fixes issue #10895.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 81 insertions, 0 deletions

diff --git a/tests/basics/assign_expr_scope.py b/tests/basics/assign_expr_scope.py
new file mode 100644
index 000000000..69dc9f0d1
--- /dev/null
+++ b/tests/basics/assign_expr_scope.py
@@ -0,0 +1,81 @@
+# Test scoping rules for assignment expression :=.
+
+# Test that var is closed over (assigned to in the scope of scope0).
+def scope0():
+    any((var0 := i) for i in range(2))
+    return var0
+
+
+print("scope0")
+print(scope0())
+print(globals().get("var0", None))
+
+# Test that var1 gets closed over correctly in the list comprehension.
+def scope1():
+    var1 = 0
+    dummy1 = 1
+    dummy2 = 1
+    print([var1 := i for i in [0, 1] if i > var1])
+    print(var1)
+
+
+print("scope1")
+scope1()
+print(globals().get("var1", None))
+
+# Test that var2 in the comprehension honours the global declaration.
+def scope2():
+    global var2
+    print([var2 := i for i in range(2)])
+    print(globals().get("var2", None))
+
+
+print("scope2")
+scope2()
+print(globals().get("var2", None))
+
+# Test that var1 in the comprehension remains local to inner1.
+def scope3():
+    global var3
+
+    def inner3():
+        print([var3 := i for i in range(2)])
+
+    inner3()
+    print(globals().get("var3", None))
+
+
+print("scope3")
+scope3()
+print(globals().get("var3", None))
+
+# Test that var4 in the comprehension honours the global declarations.
+def scope4():
+    global var4
+
+    def inner4():
+        global var4
+        print([var4 := i for i in range(2)])
+
+    inner4()
+    print(var4)
+
+
+print("scope4")
+scope4()
+print(globals().get("var4", None))
+
+# Test that var5 in the comprehension honours the nonlocal declaration.
+def scope5():
+    def inner5():
+        nonlocal var5
+        print([var5 := i for i in range(2)])
+
+    inner5()
+    print(var5)
+    var5 = 0  # force var5 to be a local to scope5
+
+
+print("scope5")
+scope5()
+print(globals().get("var5", None))