Add Yield() statement handler

Fix Module() handler to avoid including the doc string in the AST

1 files changed, 11 insertions, 1 deletions

diff --git a/Lib/compiler/transformer.py b/Lib/compiler/transformer.py
index 10d0e51..ebb3481 100644
--- a/Lib/compiler/transformer.py
+++ b/Lib/compiler/transformer.py
@@ -153,8 +153,12 @@ class Transformer:
 
     def file_input(self, nodelist):
         doc = self.get_docstring(nodelist, symbol.file_input)
+        if doc is not None:
+            i = 1
+        else:
+            i = 0
         stmts = []
-        for node in nodelist:
+        for node in nodelist[i:]:
             if node[0] != token.ENDMARKER and node[0] != token.NEWLINE:
                 self.com_append_stmt(stmts, node)
         return Module(doc, Stmt(stmts))
@@ -340,6 +344,11 @@ class Transformer:
         n.lineno = nodelist[0][2]
         return n
 
+    def yield_stmt(self, nodelist):
+        n = Yield(self.com_node(nodelist[1]))
+        n.lineno = nodelist[0][2]
+        return n
+
     def raise_stmt(self, nodelist):
         # raise: [test [',' test [',' test]]]
         if len(nodelist) > 5:
@@ -1245,6 +1254,7 @@ _legal_node_types = [
     symbol.continue_stmt,
     symbol.return_stmt,
     symbol.raise_stmt,
+    symbol.yield_stmt,
     symbol.import_stmt,
     symbol.global_stmt,
     symbol.exec_stmt,