compile: make decorators for functions work

ncw · ncw · commit 99c3f74e8797 · 2015-04-30T18:11:48.000+01:00
diff --git a/compile/compile.go b/compile/compile.go
@@ -459,10 +459,21 @@ func (c *compiler) compileFunc(Ast ast.Ast, Args *ast.Arguments, DecoratorList [
 		c.LoadConst(annotations)
 	}
 
+	// Load decorators onto stack
+	for _, expr := range DecoratorList {
+		c.Expr(expr)
+	}
+
+	// Make function or closure, leaving it on the stack
 	posdefaults := uint32(len(Args.Defaults))
 	kwdefaults := uint32(len(Args.KwDefaults))
 	args := uint32(posdefaults + (kwdefaults << 8) + (num_annotations << 16))
 	c.makeClosure(code, args, newC)
+
+	// Call decorators
+	for _ = range DecoratorList {
+		c.OpArg(vm.CALL_FUNCTION, 1) // 1 positional, 0 keyword pair
+	}
 }
 
 // Compile statement
diff --git a/compile/compile_data_test.go b/compile/compile_data_test.go
@@ -2576,4 +2576,119 @@ var compileTestData = []struct {
 		Firstlineno: 1,
 		Lnotab:      "",
 	}, nil, ""},
+	{"@wrap\ndef fn(o):\n    return o", "exec", &py.Code{
+		Argcount:       0,
+		Kwonlyargcount: 0,
+		Nlocals:        0,
+		Stacksize:      3,
+		Flags:          64,
+		Code:           "\x65\x00\x00\x64\x00\x00\x64\x01\x00\x84\x00\x00\x83\x01\x00\x5a\x01\x00\x64\x02\x00\x53",
+		Consts: []py.Object{&py.Code{
+			Argcount:       1,
+			Kwonlyargcount: 0,
+			Nlocals:        1,
+			Stacksize:      1,
+			Flags:          67,
+			Code:           "\x7c\x00\x00\x53",
+			Consts:         []py.Object{py.None},
+			Names:          []string{},
+			Varnames:       []string{"o"},
+			Freevars:       []string{},
+			Cellvars:       []string{},
+			Filename:       "<string>",
+			Name:           "fn",
+			Firstlineno:    1,
+			Lnotab:         "\x00\x02",
+		}, py.String("fn"), py.None},
+		Names:       []string{"wrap", "fn"},
+		Varnames:    []string{},
+		Freevars:    []string{},
+		Cellvars:    []string{},
+		Filename:    "<string>",
+		Name:        "<module>",
+		Firstlineno: 1,
+		Lnotab:      "",
+	}, nil, ""},
+	{"@wrap1\n@wrap2(\"potato\", 2)\n@wrap3(\"sausage\")\n@wrap4\ndef fn(o):\n    return o", "exec", &py.Code{
+		Argcount:       0,
+		Kwonlyargcount: 0,
+		Nlocals:        0,
+		Stacksize:      6,
+		Flags:          64,
+		Code:           "\x65\x00\x00\x65\x01\x00\x64\x00\x00\x64\x01\x00\x83\x02\x00\x65\x02\x00\x64\x02\x00\x83\x01\x00\x65\x03\x00\x64\x03\x00\x64\x04\x00\x84\x00\x00\x83\x01\x00\x83\x01\x00\x83\x01\x00\x83\x01\x00\x5a\x04\x00\x64\x05\x00\x53",
+		Consts: []py.Object{py.String("potato"), py.Int(2), py.String("sausage"), &py.Code{
+			Argcount:       1,
+			Kwonlyargcount: 0,
+			Nlocals:        1,
+			Stacksize:      1,
+			Flags:          67,
+			Code:           "\x7c\x00\x00\x53",
+			Consts:         []py.Object{py.None},
+			Names:          []string{},
+			Varnames:       []string{"o"},
+			Freevars:       []string{},
+			Cellvars:       []string{},
+			Filename:       "<string>",
+			Name:           "fn",
+			Firstlineno:    1,
+			Lnotab:         "\x00\x05",
+		}, py.String("fn"), py.None},
+		Names:       []string{"wrap1", "wrap2", "wrap3", "wrap4", "fn"},
+		Varnames:    []string{},
+		Freevars:    []string{},
+		Cellvars:    []string{},
+		Filename:    "<string>",
+		Name:        "<module>",
+		Firstlineno: 1,
+		Lnotab:      "\x03\x01\x0c\x01\x09\x01",
+	}, nil, ""},
+	{"def outer(o):\n    @wrap1\n    @wrap2(\"potato\", o)\n    def inner(i):\n        return o+i", "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:       1,
+			Kwonlyargcount: 0,
+			Nlocals:        2,
+			Stacksize:      5,
+			Flags:          3,
+			Code:           "\x74\x00\x00\x74\x01\x00\x64\x01\x00\x88\x00\x00\x83\x02\x00\x87\x00\x00\x66\x01\x00\x64\x02\x00\x64\x03\x00\x86\x00\x00\x83\x01\x00\x83\x01\x00\x7d\x01\x00\x64\x00\x00\x53",
+			Consts: []py.Object{py.None, py.String("potato"), &py.Code{
+				Argcount:       1,
+				Kwonlyargcount: 0,
+				Nlocals:        1,
+				Stacksize:      2,
+				Flags:          19,
+				Code:           "\x88\x00\x00\x7c\x00\x00\x17\x53",
+				Consts:         []py.Object{py.None},
+				Names:          []string{},
+				Varnames:       []string{"i"},
+				Freevars:       []string{"o"},
+				Cellvars:       []string{},
+				Filename:       "<string>",
+				Name:           "inner",
+				Firstlineno:    2,
+				Lnotab:         "\x00\x03",
+			}, py.String("outer.<locals>.inner")},
+			Names:       []string{"wrap1", "wrap2"},
+			Varnames:    []string{"o", "inner"},
+			Freevars:    []string{},
+			Cellvars:    []string{"o"},
+			Filename:    "<string>",
+			Name:        "outer",
+			Firstlineno: 1,
+			Lnotab:      "\x00\x01\x03\x01",
+		}, py.String("outer"), py.None},
+		Names:       []string{"outer"},
+		Varnames:    []string{},
+		Freevars:    []string{},
+		Cellvars:    []string{},
+		Filename:    "<string>",
+		Name:        "<module>",
+		Firstlineno: 1,
+		Lnotab:      "",
+	}, nil, ""},
 }
diff --git a/compile/make_compile_test.py b/compile/make_compile_test.py
@@ -193,6 +193,23 @@ def inner2(s):
 def outer(o):
     x = 17
     return lambda a,b=42,*args,**kw: a*b*args*kw*x*o''', "exec"),
+    ('''\
+@wrap
+def fn(o):
+    return o''', "exec"),
+    ('''\
+@wrap1
+@wrap2("potato", 2)
+@wrap3("sausage")
+@wrap4
+def fn(o):
+    return o''', "exec"),
+    ('''\
+def outer(o):
+    @wrap1
+    @wrap2("potato", o)
+    def inner(i):
+        return o+i''', "exec"),
 ]
 
 def string(s):
