compile: Fix Varnames

Author: ncw

compile/compile.go

@@ -223,7 +223,7 @@ func (c *compiler) Index(Id string, Names *[]string) uint32 {
 		}
 	}
 	*Names = append(*Names, Id)
-	return uint32(len(c.Code.Names) - 1)
+	return uint32(len(*Names) - 1)
 }
 
 // Compiles a python name
@@ -307,18 +307,18 @@ func (c *compiler) Stmt(stmt ast.Stmt) {
 
 		// Arguments
 		for _, arg := range node.Args.Args {
-			code.Varnames = append(code.Varnames, string(arg.Arg))
+			c.Index(string(arg.Arg), &code.Varnames)
 		}
 		for _, arg := range node.Args.Kwonlyargs {
-			code.Varnames = append(code.Varnames, string(arg.Arg))
+			c.Index(string(arg.Arg), &code.Varnames)
 		}
 		if node.Args.Vararg != nil {
 			code.Nlocals++
-			code.Varnames = append(code.Varnames, string(node.Args.Vararg.Arg))
+			c.Index(string(node.Args.Vararg.Arg), &code.Varnames)
 		}
 		if node.Args.Kwarg != nil {
 			code.Nlocals++
-			code.Varnames = append(code.Varnames, string(node.Args.Kwarg.Arg))
+			c.Index(string(node.Args.Kwarg.Arg), &code.Varnames)
 			code.Flags |= py.CO_VARKEYWORDS
 		}
 

compile/compile_data_test.go

@@ -2182,4 +2182,37 @@ var compileTestData = []struct {
 		Firstlineno: 1,
 		Lnotab:      "",
 	}, nil, ""},
+	{"def fn(a,b): a+b+c+d", "exec", &py.Code{
+		Argcount:       0,
+		Kwonlyargcount: 0,
+		Nlocals:        0,
+		Stacksize:      2,
+		Flags:          64,
+		Code:           "\x64\x00\x00\x64\x01\x00\x84\x00\x00\x5a\x00\x00\x64\x02\x00\x53",
+		Consts: []py.Object{&py.Code{
+			Argcount:       2,
+			Kwonlyargcount: 0,
+			Nlocals:        2,
+			Stacksize:      2,
+			Flags:          67,
+			Code:           "\x7c\x00\x00\x7c\x01\x00\x17\x74\x00\x00\x17\x74\x01\x00\x17\x01\x64\x00\x00\x53",
+			Consts:         []py.Object{py.None},
+			Names:          []string{"c", "d"},
+			Varnames:       []string{"a", "b"},
+			Freevars:       []string{},
+			Cellvars:       []string{},
+			Filename:       "<string>",
+			Name:           "fn",
+			Firstlineno:    1,
+			Lnotab:         "",
+		}, py.String("fn"), py.None},
+		Names:       []string{"fn"},
+		Varnames:    []string{},
+		Freevars:    []string{},
+		Cellvars:    []string{},
+		Filename:    "<string>",
+		Name:        "<module>",
+		Firstlineno: 1,
+		Lnotab:      "",
+	}, nil, ""},
 }

compile/make_compile_test.py

@@ -164,7 +164,7 @@
     ('''def fn(a,*arg,b=1,c=2,**kwargs): pass''', "exec"),
     ('''def fn(a:"a",*arg:"arg",b:"b"=1,c:"c"=2,**kwargs:"kw") -> "ret": pass''', "exec"),
     ('''def fn(): a+b''', "exec"),
-    #('''def fn(a,b): a+b+c+d''', "exec"),
+    ('''def fn(a,b): a+b+c+d''', "exec"),
 
 ]