support for inlining the __enter__() method call in with statements

Cython/Compiler/Nodes.py

@@ -5144,21 +5144,26 @@ class WithStatNode(StatNode):
     #  manager          The with statement manager object
     #  target           ExprNode  the target lhs of the __enter__() call
     #  body             StatNode
+    #  enter_call       ExprNode  the call to the __enter__() method
 
-    child_attrs = ["manager", "target", "body"]
+    child_attrs = ["manager", "target", "body", "enter_call"]
 
+    enter_call = None
     has_target = False
 
     def analyse_declarations(self, env):
         self.manager.analyse_declarations(env)
+        self.enter_call.analyse_declarations(env)
         self.body.analyse_declarations(env)
 
     def analyse_expressions(self, env):
         self.manager.analyse_types(env)
+        self.enter_call.analyse_types(env)
         self.body.analyse_expressions(env)
 
     def generate_function_definitions(self, env, code):
         self.manager.generate_function_definitions(env, code)
+        self.enter_call.generate_function_definitions(env, code)
         self.body.generate_function_definitions(env, code)
 
     def generate_execution_code(self, code):
@@ -5177,40 +5182,23 @@ class WithStatNode(StatNode):
         old_error_label = code.new_error_label()
         intermediate_error_label = code.error_label
 
-        enter_func = code.funcstate.allocate_temp(py_object_type, manage_ref=True)
-        code.putln("%s = PyObject_GetAttr(%s, %s); %s" % (
-            enter_func,
-            self.manager.py_result(),
-            code.get_py_string_const(EncodedString('__enter__'), identifier=True),
-            code.error_goto_if_null(enter_func, self.pos),
-            ))
-        code.put_gotref(enter_func)
+        self.enter_call.generate_evaluation_code(code)
+        if not self.target:
+            self.enter_call.generate_disposal_code(code)
+            self.enter_call.free_temps(code)
         self.manager.generate_disposal_code(code)
         self.manager.free_temps(code)
-        self.target_temp.allocate(code)
-        code.putln('%s = PyObject_Call(%s, ((PyObject *)%s), NULL); %s' % (
-            self.target_temp.result(),
-            enter_func,
-            Naming.empty_tuple,
-            code.error_goto_if_null(self.target_temp.result(), self.pos),
-            ))
-        code.put_gotref(self.target_temp.result())
-        code.put_decref_clear(enter_func, py_object_type)
-        code.funcstate.release_temp(enter_func)
-        if not self.has_target:
-            code.put_decref_clear(self.target_temp.result(), type=py_object_type)
-            self.target_temp.release(code)
-            # otherwise, WithTargetAssignmentStatNode will do it for us
 
         code.error_label = old_error_label
         self.body.generate_execution_code(code)
 
-        step_over_label = code.new_label()
-        code.put_goto(step_over_label)
-        code.put_label(intermediate_error_label)
-        code.put_decref_clear(self.exit_var, py_object_type)
-        code.put_goto(old_error_label)
-        code.put_label(step_over_label)
+        if code.label_used(intermediate_error_label):
+            step_over_label = code.new_label()
+            code.put_goto(step_over_label)
+            code.put_label(intermediate_error_label)
+            code.put_decref_clear(self.exit_var, py_object_type)
+            code.put_goto(old_error_label)
+            code.put_label(step_over_label)
 
         code.funcstate.release_temp(self.exit_var)
         code.putln('}')
@@ -5222,8 +5210,9 @@ class WithTargetAssignmentStatNode(AssignmentNode):
     # This is a special cased assignment that steals the RHS reference
     # and frees its temp.
     #
-    # lhs  ExprNode  the assignment target
-    # rhs  TempNode  the return value of the __enter__() call
+    # lhs       ExprNode  the assignment target
+    # orig_rhs  ExprNode  the return value of the __enter__() call (not owned by this node!)
+    # rhs       ResultRefNode  a ResultRefNode for the orig_rhs (owned by this node)
 
     child_attrs = ["lhs", "rhs"]
 
@@ -5234,16 +5223,13 @@ class WithTargetAssignmentStatNode(AssignmentNode):
         self.rhs.analyse_types(env)
         self.lhs.analyse_target_types(env)
         self.lhs.gil_assignment_check(env)
-        self.orig_rhs = self.rhs
         self.rhs = self.rhs.coerce_to(self.lhs.type, env)
 
     def generate_execution_code(self, code):
         self.rhs.generate_evaluation_code(code)
         self.lhs.generate_assignment_code(self.rhs, code)
-        self.orig_rhs.release(code)
-
-    def generate_function_definitions(self, env, code):
-        self.rhs.generate_function_definitions(env, code)
+        self.orig_rhs.generate_disposal_code(code)
+        self.orig_rhs.free_temps(code)
 
     def annotate(self, code):
         self.lhs.annotate(code)

Cython/Compiler/ParseTreeTransforms.py

@@ -1164,16 +1164,20 @@ class WithTransform(CythonTransform, SkipDeclarations):
         self.visitchildren(node, 'body')
         pos = node.pos
         body, target, manager = node.body, node.target, node.manager
-        node.target_temp = ExprNodes.TempNode(pos, type=PyrexTypes.py_object_type)
+        node.enter_call = ExprNodes.SimpleCallNode(
+            pos, function = ExprNodes.AttributeNode(
+                pos, obj = ResultRefNode(manager),
+                attribute = EncodedString('__enter__')),
+            args = [],
+            is_temp = True)
         if target is not None:
-            node.has_target = True
             body = Nodes.StatListNode(
                 pos, stats = [
                     Nodes.WithTargetAssignmentStatNode(
-                        pos, lhs = target, rhs = node.target_temp),
-                    body
-                    ])
-            node.target = None
+                        pos, lhs = target,
+                        rhs = ResultRefNode(node.enter_call),
+                        orig_rhs = node.enter_call),
+                    body])
 
         excinfo_target = ResultRefNode(
             pos=pos, type=Builtin.tuple_type, may_hold_none=False)

tests/run/type_inference.pyx

@@ -539,12 +539,47 @@ def with_statement():
     """
     >>> with_statement()
     Python object
-    'Python object'
+    Python object
     """
     x = 1.0
     with EmptyContextManager() as x:
         print(typeof(x))
-    return typeof(x)
+    print(typeof(x))
+    return x
+
+@cython.final
+cdef class TypedContextManager(object):
+    cpdef double __enter__(self):
+        return 2.0
+    def __exit__(self, *args):
+        return 0
+
+def with_statement_typed():
+    """
+    >>> with_statement_typed()
+    double
+    double
+    2.0
+    """
+    x = 1.0
+    with TypedContextManager() as x:
+        print(typeof(x))
+    print(typeof(x))
+    return x
+
+def with_statement_untyped():
+    """
+    >>> with_statement_untyped()
+    Python object
+    Python object
+    2.0
+    """
+    x = 1.0
+    cdef object t = TypedContextManager()
+    with t as x:
+        print(typeof(x))
+    print(typeof(x))
+    return x
 
 
 # Regression test for trac #638.