Use special method lookup in async for

Author: slozier

src/core/IronPython/Compiler/Ast/AsyncForStatement.cs

@@ -6,7 +6,8 @@
 
 using System.Threading;
 
-using Microsoft.Scripting;
+using IronPython.Runtime.Binding;
+
 using MSAst = System.Linq.Expressions;
 
 namespace IronPython.Compiler.Ast {
@@ -78,17 +79,15 @@ T SetScope<T>(T node) where T : Node {
             }
 
             // _iter = ITER.__aiter__()
-            var aiterAttr = SetScope(new MemberExpression(List, "__aiter__"));
-            var aiterCall = SetScope(new CallExpression(aiterAttr, null, null));
+            var aiterCall = SetScope(new UnaryExpression(PythonOperationKind.AIter, List));
             var assignIter = SetScope(new AssignmentStatement([SetScope(new NameExpression(iterName))], aiterCall));
 
             // running = True
             var trueConst = SetScope(new ConstantExpression(true));
             var assignRunning = SetScope(new AssignmentStatement([SetScope(new NameExpression(runningName))], trueConst));
 
             // TARGET = await __aiter.__anext__()
-            var anextAttr = SetScope(new MemberExpression(SetScope(new NameExpression(iterName)), "__anext__"));
-            var anextCall = SetScope(new CallExpression(anextAttr, null, null));
+            var anextCall = SetScope(new UnaryExpression(PythonOperationKind.ANext, SetScope(new NameExpression(iterName))));
             var awaitNext = new AwaitExpression(anextCall);
             var assignTarget = SetScope(new AssignmentStatement([Left], awaitNext));
 

src/core/IronPython/Compiler/Ast/UnaryExpression.cs

@@ -2,20 +2,25 @@
 // The .NET Foundation licenses this file to you under the Apache 2.0 License.
 // See the LICENSE file in the project root for more information.
 
-using MSAst = System.Linq.Expressions;
+#nullable enable
 
-using System;
 using System.Diagnostics;
 
 using IronPython.Runtime.Binding;
 
-namespace IronPython.Compiler.Ast {
-    using Ast = MSAst.Expression;
-    using AstUtils = Microsoft.Scripting.Ast.Utils;
+using MSAst = System.Linq.Expressions;
 
+namespace IronPython.Compiler.Ast {
     public class UnaryExpression : Expression {
         public UnaryExpression(PythonOperator op, Expression expression) {
             Operator = op;
+            OperationKind = PythonOperatorToOperatorString(op);
+            Expression = expression;
+            EndIndex = expression.EndIndex;
+        }
+
+        internal UnaryExpression(PythonOperationKind op, Expression expression) {
+            OperationKind = op;
             Expression = expression;
             EndIndex = expression.EndIndex;
         }
@@ -24,13 +29,10 @@ public UnaryExpression(PythonOperator op, Expression expression) {
 
         public PythonOperator Operator { get; }
 
-        public override MSAst.Expression Reduce() {
-            return GlobalParent.Operation(
-                typeof(object),
-                PythonOperatorToOperatorString(Operator),
-                Expression
-            );
-        }
+        internal PythonOperationKind OperationKind { get; }
+
+        public override MSAst.Expression Reduce()
+            => GlobalParent.Operation(typeof(object), OperationKind, Expression);
 
         public override void Walk(PythonWalker walker) {
             if (walker.Walk(this)) {

src/core/IronPython/Runtime/Binding/PythonOperationKind.cs

@@ -106,6 +106,9 @@ internal enum PythonOperationKind {
         /// </summary>
         GetEnumeratorForIteration,
 
+        AIter,
+        ANext,
+
         ///<summary>Operator for performing add</summary>
         Add,
         ///<summary>Operator for performing sub</summary>

src/core/IronPython/Runtime/Binding/PythonProtocol.Operations.cs

@@ -190,6 +190,12 @@ internal static partial class PythonProtocol {
                 case PythonOperationKind.GetEnumeratorForIteration:
                     res = MakeEnumeratorOperation(operation, args[0]);
                     break;
+                case PythonOperationKind.AIter:
+                    res = MakeUnaryOperation(operation, args[0], "__aiter__", TypeError(operation, "'async for' requires an object with __aiter__ method, got {0}", args));
+                    break;
+                case PythonOperationKind.ANext:
+                    res = MakeUnaryOperation(operation, args[0], "__anext__", TypeError(operation, "'async for' received an invalid object from __aiter__: {0}", args));
+                    break;
                 default:
                     res = BindingHelpers.AddPythonBoxing(MakeBinaryOperation(operation, args, operation.Operation, null));
                     break;

tests/suite/test_async.py

@@ -296,6 +296,20 @@ async def test():
 
         self.assertEqual(run_coro(test()), [110, 120, 210, 220])
 
+    def test_special_method_lookup(self):
+        """Ensure async for looks up __aiter__ on the type, not the instance."""
+
+        a = AsyncIter([1, 2, 3])
+        a.__aiter__ = lambda: AsyncIter([98])  # should be ignored
+        a.__anext__ = lambda: 99  # should be ignored
+
+        async def test():
+            result = []
+            async for x in a:
+                result.append(x)
+            return result
+
+        self.assertEqual(run_coro(test()), [1, 2, 3])
 
 class AsyncCombinedTest(unittest.TestCase):
     """Tests combining async with and async for."""