From 30381b3b10273f6d4f84e11b7353819065175d7e Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 26 Jul 2016 19:16:18 +0300
Subject: [PATCH] PY-19775 Fixed: Unresolved attribute reference on attribute
 defined inside async def __init__

Highlight builtin methods (except __aiter__, __aenter__, __aexit__, __anext__) which are not allowed to be async
---
 .../src/com/jetbrains/python/psi/PyFunction.java    |  4 +++-
 .../jetbrains/python/psi/impl/PyFunctionImpl.java   | 13 +++++++++++++
 .../validation/DumbAwareHighlightingAnnotator.java  | 13 ++++++++++++-
 python/testData/highlighting/asyncBuiltinMethods.py |  9 +++++++++
 .../jetbrains/python/PythonHighlightingTest.java    |  7 ++++++-
 5 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 python/testData/highlighting/asyncBuiltinMethods.py

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
index f49001b62674..8fc44672dab8 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyFunction.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2014 JetBrains s.r.o.
+ * Copyright 2000-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -70,6 +70,8 @@ public interface PyFunction extends PsiNamedElement, StubBasedPsiElement<PyFunct
 
   boolean isAsync();
 
+  boolean isAsyncAllowed();
+
   /**
    * Flags that mark common alterations of a function: decoration by and wrapping in classmethod() and staticmethod().
    */
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 9880d1789dda..f41f4104cfa3 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -31,6 +31,7 @@ import com.intellij.psi.search.SearchScope;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.*;
+import com.intellij.util.ArrayUtil;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.PlatformIcons;
 import com.jetbrains.python.PyElementTypes;
@@ -675,6 +676,18 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
     return getNode().findChildByType(PyTokenTypes.ASYNC_KEYWORD) != null;
   }
 
+  @Override
+  public boolean isAsyncAllowed() {
+    final LanguageLevel languageLevel = LanguageLevel.forElement(this);
+    final String functionName = getName();
+
+    return languageLevel.isAtLeast(LanguageLevel.PYTHON35) && (
+      functionName == null ||
+      ArrayUtil.contains(functionName, PyNames.AITER, "__anext__", "__aenter__", "__aexit__") ||
+      !PyNames.getBuiltinMethods(languageLevel).containsKey(functionName)
+    );
+  }
+
   @Nullable
   private Modifier getWrappersFromStub() {
     final StubElement parentStub = getStub().getParentStub();
diff --git a/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java b/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java
index a9e14a620a5d..7ab06c19b322 100644
--- a/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java
+++ b/python/src/com/jetbrains/python/validation/DumbAwareHighlightingAnnotator.java
@@ -26,13 +26,24 @@ import com.jetbrains.python.highlighting.PyHighlighter;
 import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.Optional;
+
 /**
  * @author vlan
  */
 public class DumbAwareHighlightingAnnotator extends PyAnnotator implements HighlightRangeExtension {
+
   @Override
   public void visitPyFunction(PyFunction node) {
-    highlightKeyword(node, PyTokenTypes.ASYNC_KEYWORD);
+    if (node.isAsyncAllowed()) {
+      highlightKeyword(node, PyTokenTypes.ASYNC_KEYWORD);
+    }
+    else {
+      Optional
+        .ofNullable(node.getNode())
+        .map(astNode -> astNode.findChildByType(PyTokenTypes.ASYNC_KEYWORD))
+        .ifPresent(asyncNode -> getHolder().createErrorAnnotation(asyncNode, "function \"" + node.getName() + "\" cannot be async"));
+    }
   }
 
   @Override
diff --git a/python/testData/highlighting/asyncBuiltinMethods.py b/python/testData/highlighting/asyncBuiltinMethods.py
new file mode 100644
index 000000000000..6a6c6123f256
--- /dev/null
+++ b/python/testData/highlighting/asyncBuiltinMethods.py
@@ -0,0 +1,9 @@
+class A:
+    <error descr="method \"__init__\" cannot be async">async</error> def __init__(self):
+        pass
+
+    <error descr="method \"__contains__\" cannot be async">async</error> def __contains__(self, value):
+        pass
+
+    async def __aiter__(self):
+        pass
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java b/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
index d61f10407586..634bc91823a7 100644
--- a/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonHighlightingTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2000-2013 JetBrains s.r.o.
+ * Copyright 2000-2016 JetBrains s.r.o.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -266,6 +266,11 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.PYTHON35, false, false);
   }
 
+  // PY-19775
+  public void testAsyncBuiltinMethods() {
+    doTest(LanguageLevel.PYTHON35, true, false);
+  }
+
   // ---
   private void doTest(final LanguageLevel languageLevel, final boolean checkWarnings, final boolean checkInfos) {
     PythonLanguageLevelPusher.setForcedLanguageLevel(myFixture.getProject(), languageLevel);