compile: finish lambdas including closures

ncw · ncw · commit 7686a574e6da · 2015-04-30T17:59:30.000+01:00
diff --git a/compile/compile.go b/compile/compile.go
@@ -143,12 +143,12 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
 		valueOnStack = true
 	case *ast.Suite:
 		c.Stmts(node.Body)
-	case ast.Expr:
+	case *ast.Lambda:
 		// Make None the first constant as lambda can't have a docstring
 		c.Const(py.None)
 		code.Name = "<lambda>"
 		c.setQualname() // FIXME is this in the right place!
-		c.Expr(node)
+		c.Expr(node.Body)
 		valueOnStack = true
 	case *ast.FunctionDef:
 		code.Name = string(node.Name)
@@ -406,6 +406,65 @@ func (c *compiler) setQualname() {
 	}
 }
 
+// Compile a function
+func (c *compiler) compileFunc(Ast ast.Ast, Args *ast.Arguments, DecoratorList []ast.Expr, Returns ast.Expr) {
+	newSymTable := c.SymTable.FindChild(Ast)
+	if newSymTable == nil {
+		panic("No symtable found for function")
+	}
+	newC := newCompiler(c, compilerScopeFunction)
+	code, err := newC.compileAst(Ast, c.Code.Filename, 0, false, newSymTable)
+	if err != nil {
+		panic(err)
+	}
+	// FIXME need these set in code before we compile - (pass in node?)
+	code.Argcount = int32(len(Args.Args))
+	code.Kwonlyargcount = int32(len(Args.Kwonlyargs))
+
+	// Defaults
+	for _, expr := range Args.Defaults {
+		c.Expr(expr)
+	}
+
+	// KwDefaults
+	if len(Args.Kwonlyargs) != len(Args.KwDefaults) {
+		panic("differing number of Kwonlyargs to KwDefaults")
+	}
+	for i := range Args.KwDefaults {
+		c.LoadConst(py.String(Args.Kwonlyargs[i].Arg))
+		c.Expr(Args.KwDefaults[i])
+	}
+
+	// Annotations
+	annotations := py.Tuple{}
+	addAnnotation := func(args ...*ast.Arg) {
+		for _, arg := range args {
+			if arg != nil && arg.Annotation != nil {
+				c.Expr(arg.Annotation)
+				annotations = append(annotations, py.String(arg.Arg))
+			}
+		}
+	}
+	addAnnotation(Args.Args...)
+	addAnnotation(Args.Vararg)
+	addAnnotation(Args.Kwonlyargs...)
+	addAnnotation(Args.Kwarg)
+	if Returns != nil {
+		c.Expr(Returns)
+		annotations = append(annotations, py.String("return"))
+	}
+	num_annotations := uint32(len(annotations))
+	if num_annotations > 0 {
+		num_annotations++ // include the tuple
+		c.LoadConst(annotations)
+	}
+
+	posdefaults := uint32(len(Args.Defaults))
+	kwdefaults := uint32(len(Args.KwDefaults))
+	args := uint32(posdefaults + (kwdefaults << 8) + (num_annotations << 16))
+	c.makeClosure(code, args, newC)
+}
+
 // Compile statement
 func (c *compiler) Stmt(stmt ast.Stmt) {
 	switch node := stmt.(type) {
@@ -415,62 +474,7 @@ func (c *compiler) Stmt(stmt ast.Stmt) {
 		// Body          []Stmt
 		// DecoratorList []Expr
 		// Returns       Expr
-		newSymTable := c.SymTable.FindChild(stmt)
-		if newSymTable == nil {
-			panic("No symtable found for function")
-		}
-		newC := newCompiler(c, compilerScopeFunction)
-		code, err := newC.compileAst(node, c.Code.Filename, 0, false, newSymTable)
-		if err != nil {
-			panic(err)
-		}
-		// FIXME need these set in code before we compile - (pass in node?)
-		code.Argcount = int32(len(node.Args.Args))
-		code.Name = string(node.Name)
-		code.Kwonlyargcount = int32(len(node.Args.Kwonlyargs))
-
-		// Defaults
-		posdefaults := uint32(len(node.Args.Defaults))
-		for _, expr := range node.Args.Defaults {
-			c.Expr(expr)
-		}
-
-		// KwDefaults
-		if len(node.Args.Kwonlyargs) != len(node.Args.KwDefaults) {
-			panic("differing number of Kwonlyargs to KwDefaults")
-		}
-		kwdefaults := uint32(len(node.Args.KwDefaults))
-		for i := range node.Args.KwDefaults {
-			c.LoadConst(py.String(node.Args.Kwonlyargs[i].Arg))
-			c.Expr(node.Args.KwDefaults[i])
-		}
-
-		// Annotations
-		annotations := py.Tuple{}
-		addAnnotation := func(args ...*ast.Arg) {
-			for _, arg := range args {
-				if arg != nil && arg.Annotation != nil {
-					c.Expr(arg.Annotation)
-					annotations = append(annotations, py.String(arg.Arg))
-				}
-			}
-		}
-		addAnnotation(node.Args.Args...)
-		addAnnotation(node.Args.Vararg)
-		addAnnotation(node.Args.Kwonlyargs...)
-		addAnnotation(node.Args.Kwarg)
-		if node.Returns != nil {
-			c.Expr(node.Returns)
-			annotations = append(annotations, py.String("return"))
-		}
-		num_annotations := uint32(len(annotations))
-		if num_annotations > 0 {
-			num_annotations++ // include the tuple
-			c.LoadConst(annotations)
-		}
-
-		args := uint32(posdefaults + (kwdefaults << 8) + (num_annotations << 16))
-		c.makeClosure(code, args, newC)
+		c.compileFunc(stmt, node.Args, node.DecoratorList, node.Returns)
 		c.NameOp(string(node.Name), ast.Store)
 
 	case *ast.ClassDef:
@@ -899,19 +903,7 @@ func (c *compiler) Expr(expr ast.Expr) {
 		// Args *Arguments
 		// Body Expr
 		// newC := Compiler
-		newSymTable := c.SymTable.FindChild(expr)
-		if newSymTable == nil {
-			panic("No symtable found for lambda")
-		}
-		newC := newCompiler(c, compilerScopeLambda)
-		code, err := newC.compileAst(node.Body, c.Code.Filename, 0, false, newSymTable)
-		if err != nil {
-			panic(err)
-		}
-
-		code.Argcount = int32(len(node.Args.Args))
-		// FIXME node.Args - more work on lambda needed
-		c.makeClosure(code, 0, newC)
+		c.compileFunc(expr, node.Args, nil, nil)
 	case *ast.IfExp:
 		// Test   Expr
 		// Body   Expr
diff --git a/compile/compile_data_test.go b/compile/compile_data_test.go
@@ -1184,6 +1184,39 @@ var compileTestData = []struct {
 		Firstlineno: 1,
 		Lnotab:      "",
 	}, nil, ""},
+	{"lambda a,b=42,*args,**kw: a*b*args*kw", "eval", &py.Code{
+		Argcount:       0,
+		Kwonlyargcount: 0,
+		Nlocals:        0,
+		Stacksize:      3,
+		Flags:          64,
+		Code:           "\x64\x00\x00\x64\x01\x00\x64\x02\x00\x84\x01\x00\x53",
+		Consts: []py.Object{py.Int(42), &py.Code{
+			Argcount:       2,
+			Kwonlyargcount: 0,
+			Nlocals:        4,
+			Stacksize:      2,
+			Flags:          79,
+			Code:           "\x7c\x00\x00\x7c\x01\x00\x14\x7c\x02\x00\x14\x7c\x03\x00\x14\x53",
+			Consts:         []py.Object{py.None},
+			Names:          []string{},
+			Varnames:       []string{"a", "b", "args", "kw"},
+			Freevars:       []string{},
+			Cellvars:       []string{},
+			Filename:       "<string>",
+			Name:           "<lambda>",
+			Firstlineno:    1,
+			Lnotab:         "",
+		}, py.String("<lambda>")},
+		Names:       []string{},
+		Varnames:    []string{},
+		Freevars:    []string{},
+		Cellvars:    []string{},
+		Filename:    "<string>",
+		Name:        "<module>",
+		Firstlineno: 1,
+		Lnotab:      "",
+	}, nil, ""},
 	{"pass", "exec", &py.Code{
 		Argcount:       0,
 		Kwonlyargcount: 0,
@@ -2494,7 +2527,7 @@ var compileTestData = []struct {
 		Firstlineno: 1,
 		Lnotab:      "",
 	}, nil, ""},
-	{"def outer(o):\n    a = 17\n    return lambda x: o+a+x", "exec", &py.Code{
+	{"def outer(o):\n    x = 17\n    return lambda a,b=42,*args,**kw: a*b*args*kw*x*o", "exec", &py.Code{
 		Argcount:       0,
 		Kwonlyargcount: 0,
 		Nlocals:        0,
@@ -2505,20 +2538,20 @@ var compileTestData = []struct {
 			Argcount:       1,
 			Kwonlyargcount: 0,
 			Nlocals:        1,
-			Stacksize:      3,
+			Stacksize:      4,
 			Flags:          3,
-			Code:           "\x64\x01\x00\x89\x00\x00\x87\x00\x00\x87\x01\x00\x66\x02\x00\x64\x02\x00\x64\x03\x00\x86\x00\x00\x53",
-			Consts: []py.Object{py.None, py.Int(17), &py.Code{
-				Argcount:       1,
+			Code:           "\x64\x01\x00\x89\x01\x00\x64\x02\x00\x87\x00\x00\x87\x01\x00\x66\x02\x00\x64\x03\x00\x64\x04\x00\x86\x01\x00\x53",
+			Consts: []py.Object{py.None, py.Int(17), py.Int(42), &py.Code{
+				Argcount:       2,
 				Kwonlyargcount: 0,
-				Nlocals:        1,
+				Nlocals:        4,
 				Stacksize:      2,
-				Flags:          19,
-				Code:           "\x88\x01\x00\x88\x00\x00\x17\x7c\x00\x00\x17\x53",
+				Flags:          31,
+				Code:           "\x7c\x00\x00\x7c\x01\x00\x14\x7c\x02\x00\x14\x7c\x03\x00\x14\x88\x01\x00\x14\x88\x00\x00\x14\x53",
 				Consts:         []py.Object{py.None},
 				Names:          []string{},
-				Varnames:       []string{"x"},
-				Freevars:       []string{"a", "o"},
+				Varnames:       []string{"a", "b", "args", "kw"},
+				Freevars:       []string{"o", "x"},
 				Cellvars:       []string{},
 				Filename:       "<string>",
 				Name:           "<lambda>",
@@ -2528,7 +2561,7 @@ var compileTestData = []struct {
 			Names:       []string{},
 			Varnames:    []string{"o"},
 			Freevars:    []string{},
-			Cellvars:    []string{"a", "o"},
+			Cellvars:    []string{"o", "x"},
 			Filename:    "<string>",
 			Name:        "outer",
 			Firstlineno: 1,
diff --git a/compile/make_compile_test.py b/compile/make_compile_test.py
@@ -93,7 +93,7 @@
     # lambda
     ('''lambda: 0''', "eval"),
     ('''lambda x: 2*x''', "eval"),
-    # FIXME ('''lambda x=42: 2*x''', "eval"),
+    ('''lambda a,b=42,*args,**kw: a*b*args*kw''', "eval"),
     # pass statment
     ('''pass''', "exec"),
     # expr statement
@@ -191,9 +191,8 @@ def inner2(s):
     return inner''', "exec"),
     ('''\
 def outer(o):
-    a = 17
-    return lambda x: o+a+x''', "exec"),
-
+    x = 17
+    return lambda a,b=42,*args,**kw: a*b*args*kw*x*o''', "exec"),
 ]
 
 def string(s):
