From 299a99776f6890f6a8ef80276267dbe4c0e138a5 Mon Sep 17 00:00:00 2001
From: Vitja Makarov <vitja.makarov@gmail.com>
Date: Fri, 10 Dec 2010 21:43:33 +0300
Subject: [PATCH] Add close() support

---
 Cython/Compiler/ExprNodes.py           | 26 ++++++++++++++-
 Cython/Compiler/Nodes.py               |  2 ++
 Cython/Compiler/ParseTreeTransforms.py |  6 ++--
 tests/run/generators.pyx               | 44 ++++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
index 582e34d2..fd1e66a5 100755
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -5019,7 +5019,8 @@ class YieldExprNode(ExprNode):
             code.putln('%s = %s->%s;' % (cname, Naming.cur_scope_cname, save_cname))
             if type.is_pyobject:
                 code.putln('%s->%s = 0;' % (Naming.cur_scope_cname, save_cname))
-        code.putln('%s = __pyx_send_value;' % self.result())
+        code.putln('%s = __pyx_send_value; %s' %
+                   (self.result(), code.error_goto_if_null(self.result(), self.pos)))
         code.put_incref(self.result(), py_object_type)
 
 class StopIterationNode(Node):
@@ -8306,6 +8307,7 @@ generator_utility_code = UtilityCode(
 proto="""
 static PyObject *__CyGenerator_Next(PyObject *self);
 static PyObject *__CyGenerator_Send(PyObject *self, PyObject *value);
+static PyObject *__CyGenerator_Close(PyObject *self);
 typedef PyObject *(*__cygenerator_body_t)(PyObject *, PyObject *, int);
 """,
 impl="""
@@ -8355,4 +8357,26 @@ static PyObject *__CyGenerator_Send(PyObject *self, PyObject *value)
 {
     return __CyGenerator_SendEx((struct __CyGenerator *) self, value, 0);
 }
+
+static PyObject *__CyGenerator_Close(PyObject *self)
+{
+    struct __CyGenerator *generator = (struct __CyGenerator *) self;
+    PyObject *retval;
+    PyErr_SetNone(PyExc_GeneratorExit);
+    retval = __CyGenerator_SendEx(generator, NULL, 0);
+    if (retval) {
+        Py_DECREF(retval);
+        PyErr_SetString(PyExc_RuntimeError,
+                        "generator ignored GeneratorExit");
+        return NULL;
+    }
+    if (PyErr_ExceptionMatches(PyExc_StopIteration)
+        || PyErr_ExceptionMatches(PyExc_GeneratorExit))
+    {
+        PyErr_Clear();          /* ignore these errors */
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+    return NULL;
+}
 """, proto_block='utility_code_proto_before_types')
diff --git a/Cython/Compiler/Nodes.py b/Cython/Compiler/Nodes.py
index 2abe6126..48eb131b 100644
--- a/Cython/Compiler/Nodes.py
+++ b/Cython/Compiler/Nodes.py
@@ -1362,6 +1362,8 @@ class FuncDefNode(StatNode, BlockNode):
             first_run_label = code.new_label('first_run')
             code.use_label(first_run_label)
             code.put_label(first_run_label)
+            code.putln('%s' %
+                       (code.error_goto_if_null('__pyx_send_value', self.pos)))
         if not self.is_generator:
             self.generate_argument_parsing_code(env, code)
         # If an argument is assigned to in the body, we must 
diff --git a/Cython/Compiler/ParseTreeTransforms.py b/Cython/Compiler/ParseTreeTransforms.py
index dd746a0f..f8728b4f 100644
--- a/Cython/Compiler/ParseTreeTransforms.py
+++ b/Cython/Compiler/ParseTreeTransforms.py
@@ -1475,9 +1475,9 @@ class CreateClosureClasses(CythonTransform):
         e.func_cname = '__CyGenerator_Send'
         e.signature = TypeSlots.binaryfunc
 
-        #e = klass.declare_pyfunction('close', pos)
-        #e.func_cname = '__CyGenerator_Close'
-        #e.signature = TypeSlots.unaryfunc
+        e = klass.declare_pyfunction('close', pos)
+        e.func_cname = '__CyGenerator_Close'
+        e.signature = TypeSlots.unaryfunc
 
         #e = klass.declare_pyfunction('throw', pos)
         #e.func_cname = '__CyGenerator_Throw'
diff --git a/tests/run/generators.pyx b/tests/run/generators.pyx
index b6176f00..e091b63d 100644
--- a/tests/run/generators.pyx
+++ b/tests/run/generators.pyx
@@ -9,9 +9,14 @@ def very_simple():
     >>> next(x)
     Traceback (most recent call last):
     StopIteration
+    >>> x = very_simple()
+    >>> x.send(1)
+    Traceback (most recent call last):
+    TypeError: can't send non-None value to a just-started generator
     """
     yield 1
 
+
 def simple():
     """
     >>> x = simple()
@@ -80,3 +85,42 @@ def with_outer_raising(*args):
             yield i
         raise StopIteration
     return generator
+
+def test_close():
+    """
+    >>> x = test_close()
+    >>> x.close()
+    >>> x = test_close()
+    >>> next(x)
+    >>> x.close()
+    >>> x.next()
+    Traceback (most recent call last):
+    StopIteration
+    """
+    while True:
+        yield
+
+def test_ignore_close():
+    """
+    >>> x = test_ignore_close()
+    >>> x.close()
+    >>> x = test_ignore_close()
+    >>> next(x)
+    >>> x.close()
+    Traceback (most recent call last):
+    RuntimeError: generator ignored GeneratorExit
+    """
+    try:
+        yield
+    except GeneratorExit:
+        yield
+
+class Foo(object):
+    """
+    >>> obj = Foo()
+    >>> list(obj.simple(1, 2, 3))
+    [1, 2, 3]
+    """
+    def simple(self, *args):
+        for i in args:
+            yield i
-- 
2.26.2