symtable: fix yield/yield from to set generator flag

Author: ncw

symtable/make_symtable_test.py

@@ -173,6 +173,15 @@ def method(self):
 def fn():
   from potato import *
 ''', "exec", SyntaxError),
+    # Yield
+    ('''\
+def f():
+    yield
+    ''', "exec"),
+    ('''\
+def f():
+    yield from range(10)
+    ''', "exec"),
 ]
 
 def dump_bool(b):

symtable/symtable.go

@@ -357,6 +357,8 @@ func (st *SymTable) Parse(Ast ast.Ast) {
 			if node.Value != nil {
 				st.ReturnsValue = true
 			}
+		case *ast.Yield, *ast.YieldFrom:
+			st.Generator = true
 		}
 		return true
 	})

symtable/symtable_data_test.go

@@ -2699,4 +2699,89 @@ var symtableTestData = []struct {
 		Children:          Children{},
 	}, nil, ""},
 	{"def fn():\n  from potato import *\n", "exec", nil, py.SyntaxError, "import * only allowed at module level"},
+	{"def f():\n    yield\n    ", "exec", &SymTable{
+		Type:              ModuleBlock,
+		Name:              "top",
+		Lineno:            0,
+		Unoptimized:       optTopLevel,
+		Nested:            false,
+		Free:              false,
+		ChildFree:         false,
+		Generator:         false,
+		Varargs:           false,
+		Varkeywords:       false,
+		ReturnsValue:      false,
+		NeedsClassClosure: false,
+		Varnames:          []string{},
+		Symbols: Symbols{
+			"f": Symbol{
+				Flags: DefLocal,
+				Scope: ScopeLocal,
+			},
+		},
+		Children: Children{
+			&SymTable{
+				Type:              FunctionBlock,
+				Name:              "f",
+				Lineno:            1,
+				Unoptimized:       0,
+				Nested:            false,
+				Free:              false,
+				ChildFree:         false,
+				Generator:         true,
+				Varargs:           false,
+				Varkeywords:       false,
+				ReturnsValue:      false,
+				NeedsClassClosure: false,
+				Varnames:          []string{},
+				Symbols:           Symbols{},
+				Children:          Children{},
+			},
+		},
+	}, nil, ""},
+	{"def f():\n    yield from range(10)\n    ", "exec", &SymTable{
+		Type:              ModuleBlock,
+		Name:              "top",
+		Lineno:            0,
+		Unoptimized:       optTopLevel,
+		Nested:            false,
+		Free:              false,
+		ChildFree:         false,
+		Generator:         false,
+		Varargs:           false,
+		Varkeywords:       false,
+		ReturnsValue:      false,
+		NeedsClassClosure: false,
+		Varnames:          []string{},
+		Symbols: Symbols{
+			"f": Symbol{
+				Flags: DefLocal,
+				Scope: ScopeLocal,
+			},
+		},
+		Children: Children{
+			&SymTable{
+				Type:              FunctionBlock,
+				Name:              "f",
+				Lineno:            1,
+				Unoptimized:       0,
+				Nested:            false,
+				Free:              false,
+				ChildFree:         false,
+				Generator:         true,
+				Varargs:           false,
+				Varkeywords:       false,
+				ReturnsValue:      false,
+				NeedsClassClosure: false,
+				Varnames:          []string{},
+				Symbols: Symbols{
+					"range": Symbol{
+						Flags: DefUse,
+						Scope: ScopeGlobalImplicit,
+					},
+				},
+				Children: Children{},
+			},
+		},
+	}, nil, ""},
 }