compile,py: fix closure

Author: wetor

compile/compile.go

@@ -1,4 +1,4 @@
-// Copyright 2018 The go-python Authors.  All rights reserved.
+// Copyright 2023 The go-python Authors.  All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
@@ -214,12 +214,16 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
 		panic("suite should not be possible")
 	case *ast.Lambda:
 		// Make None the first constant as lambda can't have a docstring
+		code.Argcount = int32(len(node.Args.Args))
+		code.Kwonlyargcount = int32(len(node.Args.Kwonlyargs))
 		c.Const(py.None)
 		code.Name = "<lambda>"
 		c.setQualname()
 		c.Expr(node.Body)
 		valueOnStack = true
 	case *ast.FunctionDef:
+		code.Argcount = int32(len(node.Args.Args))
+		code.Kwonlyargcount = int32(len(node.Args.Kwonlyargs))
 		code.Name = string(node.Name)
 		c.setQualname()
 		c.Stmts(c.docString(node.Body, true))
@@ -295,10 +299,23 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
 		}
 		c.Op(vm.RETURN_VALUE)
 	}
-	code.Code = c.OpCodes.Assemble()
-	code.Stacksize = int32(c.OpCodes.StackDepth())
-	code.Nlocals = int32(len(code.Varnames))
-	code.Lnotab = string(c.OpCodes.Lnotab())
+	c.Code = py.MakeCode(
+		code.Argcount,
+		code.Kwonlyargcount,
+		int32(len(code.Varnames)),
+		int32(c.OpCodes.StackDepth()),
+		code.Flags,
+		c.OpCodes.Assemble(),
+		code.Consts,
+		code.Names,
+		code.Varnames,
+		code.Freevars,
+		code.Cellvars,
+		filename,
+		code.Name,
+		code.Firstlineno,
+		string(c.OpCodes.Lnotab()),
+	)
 	return nil
 }
 

py/code.go

@@ -1,4 +1,4 @@
-// Copyright 2018 The go-python Authors.  All rights reserved.
+// Copyright 2023 The go-python Authors.  All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
@@ -112,8 +112,6 @@ func NewCode(argcount int32, kwonlyargcount int32,
 	filename_ Object, name_ Object, firstlineno int32,
 	lnotab_ Object) *Code {
 
-	var cell2arg []byte
-
 	// Type assert the objects
 	consts := consts_.(Tuple)
 	namesTuple := names_.(Tuple)
@@ -154,7 +152,6 @@ func NewCode(argcount int32, kwonlyargcount int32,
 	// 	return nil;
 	// }
 
-	n_cellvars := len(cellvars)
 	intern_strings(namesTuple)
 	intern_strings(varnamesTuple)
 	intern_strings(freevarsTuple)
@@ -167,6 +164,38 @@ func NewCode(argcount int32, kwonlyargcount int32,
 			}
 		}
 	}
+
+	return MakeCode(argcount,
+		kwonlyargcount,
+		nlocals,
+		stacksize,
+		flags,
+		code,
+		consts,
+		names,
+		varnames,
+		freevars,
+		cellvars,
+		filename,
+		name,
+		firstlineno,
+		lnotab)
+}
+
+// Make a new code object
+func MakeCode(argcount int32, kwonlyargcount int32,
+	nlocals int32, stacksize int32, flags int32,
+	code string, consts_ Object, names []string,
+	varnames []string, freevars []string, cellvars []string,
+	filename string, name string, firstlineno int32,
+	lnotab string) *Code {
+
+	var cell2arg []byte
+
+	// Type assert the objects
+	consts := consts_.(Tuple)
+
+	n_cellvars := len(cellvars)
 	/* Create mapping between cells and arguments if needed. */
 	if n_cellvars != 0 {
 		total_args := argcount + kwonlyargcount

vm/tests/class.py

@@ -47,17 +47,16 @@ def method1(self, x):
 c = x()
 assert c.method1(1) == 2
 
-# FIXME doesn't work
-# doc="CLASS_DEREF2"
-# def classderef2(x):
-#     class DeRefTest:
-#         VAR = x
-#         def method1(self, x):
-#             "method1"
-#             return self.VAR+x
-#     return DeRefTest
-# x = classderef2(1)
-# c = x()
-# assert c.method1(1) == 2
+doc="CLASS_DEREF2"
+def classderef2(x):
+    class DeRefTest:
+        VAR = x
+        def method1(self, x):
+            "method1"
+            return self.VAR+x
+    return DeRefTest
+x = classderef2(1)
+c = x()
+assert c.method1(1) == 2
 
 doc="finished"

vm/tests/functions.py

@@ -22,17 +22,16 @@ def fn2(x,y=1):
 
 # Closure
 
-# FIXME something wrong with closures over function arguments...
-# doc="counter3"
-# def counter3(x):
-#     def inc():
-#         nonlocal x
-#         x += 1
-#         return x
-#     return inc
-# fn3 = counter3(1)
-# assert fn3() == 2
-# assert fn3() == 3
+doc="counter3"
+def counter3(x):
+    def inc():
+        nonlocal x
+        x += 1
+        return x
+    return inc
+fn3 = counter3(1)
+assert fn3() == 2
+assert fn3() == 3
 
 doc="counter4"
 def counter4(initial):

vm/tests2/closure.py

@@ -0,0 +1,44 @@
+# temp file
+
+doc="closure"
+def closure(x):
+    def inc():
+        nonlocal x
+        x += 1
+        return x
+    return inc
+fn3 = closure(1)
+assert fn3() == 2
+assert fn3() == 3
+
+doc="CLASS_DEREF2"
+def classderef2(x):
+    class DeRefTest:
+        VAR = x
+        def method1(self, x):
+            "method1"
+            return self.VAR+x
+    return DeRefTest
+x = classderef2(1)
+c = x()
+assert c.method1(1) == 2
+
+doc="decorator"
+
+def decorator(func):
+    def warp(*args, **kwargs):
+        print('before called')
+        func(*args, **kwargs)
+        print('after called')
+    return warp
+
+@decorator
+def func():
+    print('my called')
+
+func()
+# before called
+# my called
+# after called
+
+doc="finished"

vm/vm_test.go

@@ -19,6 +19,11 @@ func TestVm(t *testing.T) {
 	pytest.RunTests(t, "tests")
 }
 
+func TestVm2(t *testing.T) {
+	// TODO(wetor) temp test
+	pytest.RunTests(t, "tests2")
+}
+
 func BenchmarkVM(b *testing.B) {
 	pytest.RunBenchmarks(b, "benchmarks")
 }