Merge pull request #1323 from undingen/lambda_yield

don't handle functions containing lambdas with yields as generators

src/analysis/scoping_analysis.cpp

@@ -26,31 +26,28 @@ namespace pyston {
 
 class YieldVisitor : public NoopASTVisitor {
 public:
-    YieldVisitor() : containsYield(false) {}
+    AST* starting_node;
+    bool contains_yield;
 
-    bool visit_functiondef(AST_FunctionDef*) override { return true; }
+    YieldVisitor(AST* initial_node) : starting_node(initial_node), contains_yield(false) {}
+
+    // we are only interested if the statements of the initial node contain a yield not if any child function contains a
+    // yield
+    bool shouldSkip(AST* node) const { return starting_node != node; }
+    bool visit_classdef(AST_ClassDef* node) override { return shouldSkip(node); }
+    bool visit_functiondef(AST_FunctionDef* node) override { return shouldSkip(node); }
+    bool visit_lambda(AST_Lambda* node) override { return shouldSkip(node); }
 
     bool visit_yield(AST_Yield*) override {
-        containsYield = true;
+        contains_yield = true;
         return true;
     }
-
-    bool containsYield;
 };
 
 bool containsYield(AST* ast) {
-    YieldVisitor visitor;
-    if (ast->type == AST_TYPE::FunctionDef) {
-        AST_FunctionDef* funcDef = static_cast<AST_FunctionDef*>(ast);
-        for (auto& e : funcDef->body) {
-            e->accept(&visitor);
-            if (visitor.containsYield)
-                return true;
-        }
-    } else {
+    YieldVisitor visitor(ast);
     ast->accept(&visitor);
-    }
-    return visitor.containsYield;
+    return visitor.contains_yield;
 }
 
 // TODO

src/runtime/generator.cpp

@@ -467,7 +467,7 @@ extern "C" BoxedGenerator::BoxedGenerator(BoxedFunctionBase* function, Box* arg1
         this->args = new (numArgs) GCdArray();
         memcpy(&this->args->elts[0], args, numArgs * sizeof(Box*));
         for (int i = 0; i < numArgs; i++) {
-            Py_INCREF(args[i]);
+            Py_XINCREF(args[i]);
         }
     }
 

test/tests/generators.py

@@ -134,3 +134,16 @@ x = lambda: (yield 1)
 print list(x())
 x = lambda: ((yield 1), (yield 2))
 print list(x())
+
+# we used to think that this function is a generator
+def this_is_not_generator():
+    type((lambda: (yield)))
+    type((lambda: (yield))())
+    def f():
+        yield
+print type(this_is_not_generator())
+
+# we used to crash when a generator had more than 3 arguments and kwargs (because they can be NULL)
+def G(a, b, c, *args, **kwargs):
+    yield 1
+print list(G(1,2,2))