@@ -188,6 +188,33 @@ func (c *compiler) compileAst(Ast ast.Ast, filename string, futureFlags int, don
188
188
c .LoadConst (py .None )
189
189
}
190
190
c .Op (vm .RETURN_VALUE )
191
+ case * ast.ListComp :
192
+ // Elt Expr
193
+ // Generators []Comprehension
194
+ valueOnStack = true
195
+ code .Name = "<listcomp>"
196
+ c .OpArg (vm .BUILD_LIST , 0 )
197
+ c .comprehensionGenerator (node .Generators , 0 , node .Elt , nil , Ast )
198
+ case * ast.SetComp :
199
+ // Elt Expr
200
+ // Generators []Comprehension
201
+ valueOnStack = true
202
+ code .Name = "<setcomp>"
203
+ c .OpArg (vm .BUILD_SET , 0 )
204
+ c .comprehensionGenerator (node .Generators , 0 , node .Elt , nil , Ast )
205
+ case * ast.DictComp :
206
+ // Key Expr
207
+ // Value Expr
208
+ // Generators []Comprehension
209
+ valueOnStack = true
210
+ code .Name = "<dictcomp>"
211
+ c .OpArg (vm .BUILD_MAP , 0 )
212
+ c .comprehensionGenerator (node .Generators , 0 , node .Key , node .Value , Ast )
213
+ case * ast.GeneratorExp :
214
+ // Elt Expr
215
+ // Generators []Comprehension
216
+ code .Name = "<genexpr>"
217
+ c .comprehensionGenerator (node .Generators , 0 , node .Elt , nil , Ast )
191
218
192
219
default :
193
220
panic (py .ExceptionNewf (py .SyntaxError , "Unknown ModuleBase: %v" , Ast ))
@@ -961,6 +988,97 @@ func (c *compiler) callHelper(n int, Args []ast.Expr, Keywords []*ast.Keyword, S
961
988
c .OpArg (op , uint32 (args + kwargs << 8 ))
962
989
}
963
990
991
+ /* List and set comprehensions and generator expressions work by creating a
992
+ nested function to perform the actual iteration. This means that the
993
+ iteration variables don't leak into the current scope.
994
+ The defined function is called immediately following its definition, with the
995
+ result of that call being the result of the expression.
996
+ The LC/SC version returns the populated container, while the GE version is
997
+ flagged in symtable.c as a generator, so it returns the generator object
998
+ when the function is called.
999
+ This code *knows* that the loop cannot contain break, continue, or return,
1000
+ so it cheats and skips the SETUP_LOOP/POP_BLOCK steps used in normal loops.
1001
+
1002
+ Possible cleanups:
1003
+ - iterate over the generator sequence instead of using recursion
1004
+ */
1005
+ func (c * compiler ) comprehensionGenerator (generators []ast.Comprehension , gen_index int , elt ast.Expr , val ast.Expr , Ast ast.Ast ) {
1006
+ // generate code for the iterator, then each of the ifs,
1007
+ // and then write to the element
1008
+ start := new (Label )
1009
+ skip := new (Label )
1010
+ anchor := new (Label )
1011
+ gen := generators [gen_index ]
1012
+ if gen_index == 0 {
1013
+ /* Receive outermost iter as an implicit argument */
1014
+ c .Code .Argcount = 1
1015
+ c .OpArg (vm .LOAD_FAST , 0 )
1016
+ } else {
1017
+ /* Sub-iter - calculate on the fly */
1018
+ c .Expr (gen .Iter )
1019
+ c .Op (vm .GET_ITER )
1020
+ }
1021
+ c .Label (start )
1022
+ c .Jump (vm .FOR_ITER , anchor )
1023
+ c .Expr (gen .Target )
1024
+
1025
+ /* XXX this needs to be cleaned up...a lot! */
1026
+ for _ , e := range gen .Ifs {
1027
+ c .Expr (e )
1028
+ c .Jump (vm .POP_JUMP_IF_FALSE , start )
1029
+ }
1030
+
1031
+ gen_index ++
1032
+ if gen_index < len (generators ) {
1033
+ c .comprehensionGenerator (generators , gen_index , elt , val , Ast )
1034
+ }
1035
+
1036
+ /* only append after the last for generator */
1037
+ if gen_index >= len (generators ) {
1038
+ /* comprehension specific code */
1039
+ switch Ast .(type ) {
1040
+ case * ast.GeneratorExp :
1041
+ c .Expr (elt )
1042
+ c .Op (vm .YIELD_VALUE )
1043
+ c .Op (vm .POP_TOP )
1044
+ case * ast.ListComp :
1045
+ c .Expr (elt )
1046
+ c .OpArg (vm .LIST_APPEND , uint32 (gen_index + 1 ))
1047
+ case * ast.SetComp :
1048
+ c .Expr (elt )
1049
+ c .OpArg (vm .SET_ADD , uint32 (gen_index + 1 ))
1050
+ case * ast.DictComp :
1051
+ // With 'd[k] = v', v is evaluated before k, so we do the same.
1052
+ c .Expr (val )
1053
+ c .Expr (elt )
1054
+ c .OpArg (vm .MAP_ADD , uint32 (gen_index + 1 ))
1055
+ default :
1056
+ panic (fmt .Sprintf ("unknown comprehension %v" , Ast ))
1057
+ }
1058
+ c .Label (skip )
1059
+ }
1060
+ c .Jump (vm .JUMP_ABSOLUTE , start )
1061
+ c .Label (anchor )
1062
+ }
1063
+
1064
+ // Compile a comprehension
1065
+ func (c * compiler ) comprehension (expr ast.Expr , generators []ast.Comprehension ) {
1066
+ newSymTable := c .SymTable .FindChild (expr )
1067
+ if newSymTable == nil {
1068
+ panic ("No symtable found for comprehension" )
1069
+ }
1070
+ newC := newCompiler (c , compilerScopeComprehension )
1071
+ code , err := newC .compileAst (expr , c .Code .Filename , 0 , false , newSymTable )
1072
+ if err != nil {
1073
+ panic (err )
1074
+ }
1075
+ c .makeClosure (code , 0 , newC , newC .Code .Name )
1076
+ outermost_iter := generators [0 ].Iter
1077
+ c .Expr (outermost_iter )
1078
+ c .Op (vm .GET_ITER )
1079
+ c .OpArg (vm .CALL_FUNCTION , 1 )
1080
+ }
1081
+
964
1082
// Compile expressions
965
1083
func (c * compiler ) Expr (expr ast.Expr ) {
966
1084
switch node := expr .(type ) {
@@ -1078,20 +1196,20 @@ func (c *compiler) Expr(expr ast.Expr) {
1078
1196
case * ast.ListComp :
1079
1197
// Elt Expr
1080
1198
// Generators []Comprehension
1081
- panic ( "FIXME compile: ListComp not implemented" )
1199
+ c . comprehension ( expr , node . Generators )
1082
1200
case * ast.SetComp :
1083
1201
// Elt Expr
1084
1202
// Generators []Comprehension
1085
- panic ( "FIXME compile: SetComp not implemented" )
1203
+ c . comprehension ( expr , node . Generators )
1086
1204
case * ast.DictComp :
1087
1205
// Key Expr
1088
1206
// Value Expr
1089
1207
// Generators []Comprehension
1090
- panic ( "FIXME compile: DictComp not implemented" )
1208
+ c . comprehension ( expr , node . Generators )
1091
1209
case * ast.GeneratorExp :
1092
1210
// Elt Expr
1093
1211
// Generators []Comprehension
1094
- panic ( "FIXME compile: GeneratorExp not implemented" )
1212
+ c . comprehension ( expr , node . Generators )
1095
1213
case * ast.Yield :
1096
1214
// Value Expr
1097
1215
panic ("FIXME compile: Yield not implemented" )
0 commit comments