PY-28228 Resolve annotation forward references according to PEP 563

Forward references resolution implemented for annotations according to
PEP 563. Inspections fixed to respect forward references for both
annotations and pyi stubs.

PyiReferenceResolveProvider removed since its functionality is now
implemented by PyForwardReferenceResolveProvider.

python/psi-api/src/com/jetbrains/python/psi/FutureFeature.java

@@ -31,6 +31,7 @@ public enum FutureFeature {
   PRINT_FUNCTION("print_function", 26, 30),
   UNICODE_LITERALS("unicode_literals", 26, 30),
   BARRY_AS_FLUFL("barry_as_FLUFL", 31, 39), // last as of CPython 3.2
+  ANNOTATIONS("annotations", 37, 40)
   // NOTE: only add new features to the end unless you want to break existing stubs that rely on ordinal
   ;
   // TODO: link it to LanguageLevel
@@ -87,7 +88,5 @@ public enum FutureFeature {
     return level.getVersion() >= myRequiredVersion;
   }
 
-  public static final FutureFeature[] ALL = {
-    GENERATORS, DIVISION, ABSOLUTE_IMPORT, WITH_STATEMENT, PRINT_FUNCTION, UNICODE_LITERALS, BARRY_AS_FLUFL
-  };
+  public static final FutureFeature[] ALL = FutureFeature.values();
 }

python/src/META-INF/python-core-common.xml

@@ -727,7 +727,7 @@
     <typeProvider implementation="com.jetbrains.python.pyi.PyiTypeProvider"/>
     <pyModuleMembersProvider implementation="com.jetbrains.python.pyi.PyiModuleMembersProvider"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.pyi.PyiClassMembersProvider"/>
-    <pyReferenceResolveProvider implementation="com.jetbrains.python.pyi.PyiReferenceResolveProvider"/>
+    <pyReferenceResolveProvider implementation="com.jetbrains.python.psi.resolve.PyForwardReferenceResolveProvider"/>
     <visitorFilter language="PythonStub" implementationClass="com.jetbrains.python.pyi.PyiVisitorFilter"/>
     <inspectionExtension implementation="com.jetbrains.python.pyi.PyiInspectionExtension"/>
     <inspectionExtension implementation="com.jetbrains.python.codeInsight.typing.PyTypingInspectionExtension"/>

python/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java

@@ -39,6 +39,7 @@ import com.jetbrains.python.inspections.quickfix.AddGlobalQuickFix;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyGlobalStatementNavigator;
+import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 
@@ -68,8 +69,13 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
     public Visitor(final ProblemsHolder holder, LocalInspectionToolSession session) {
       super(holder, session);
     }
+
     @Override
     public void visitPyReferenceExpression(final PyReferenceExpression node) {
+      if (PyResolveUtil.allowForwardReferences(node)) {
+        return;
+      }
+
       if (node.getContainingFile() instanceof PyExpressionCodeFragment) {
         return;
       }
@@ -127,9 +133,6 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
           return;
         }
         final PsiPolyVariantReference ref = node.getReference(getResolveContext());
-        if (ref == null) {
-          return;
-        }
         final PsiElement resolved = ref.resolve();
         final boolean isBuiltin = PyBuiltinCache.getInstance(node).isBuiltin(resolved);
         if (owner instanceof PyClass) {

python/src/com/jetbrains/python/psi/resolve/PyForwardReferenceResolveProvider.kt

@@ -0,0 +1,34 @@
+// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.jetbrains.python.psi.resolve
+
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
+import com.jetbrains.python.psi.PyQualifiedExpression
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+/**
+ * Forward references resolution for annotations and pyi stubs.
+ *
+ * @see <a href="PEP-484">https://www.python.org/dev/peps/pep-0484/</a>
+ * @see <a href="PEP-563">https://www.python.org/dev/peps/pep-0563/</a>
+ */
+class PyForwardReferenceResolveProvider : PyReferenceResolveProvider {
+
+  override fun resolveName(element: PyQualifiedExpression, context: TypeEvalContext): List<RatedResolveResult> {
+    if (!PyResolveUtil.allowForwardReferences(element)) {
+      return emptyList()
+    }
+
+    val referencedName = element.referencedName ?: return emptyList()
+    val originalOwner = ScopeUtil.getScopeOwner(element)
+
+    return if (originalOwner != null) {
+      PyResolveUtil.resolveLocally(originalOwner, referencedName)
+        .map { RatedResolveResult(RatedResolveResult.RATE_NORMAL, it) }
+    } else emptyList()
+  }
+
+  override fun allowsForwardOutgoingReferencesInClass(element: PyQualifiedExpression): Boolean {
+    return PyResolveUtil.allowForwardReferences(element)
+  }
+
+}

python/src/com/jetbrains/python/psi/resolve/PyResolveUtil.java

@@ -16,10 +16,13 @@
 package com.jetbrains.python.psi.resolve;
 
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.ResolveState;
 import com.intellij.psi.scope.PsiScopeProcessor;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
+import com.jetbrains.extenstions.PsiElementExtKt;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
@@ -30,6 +33,7 @@ import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.types.PyClassLikeType;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
+import com.jetbrains.python.pyi.PyiUtil;
 import one.util.streamex.StreamEx;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -221,4 +225,23 @@ public class PyResolveUtil {
 
     return null;
   }
+
+  /**
+   * Check whether forward references are allowed for the given element.
+   */
+  public static boolean allowForwardReferences(@NotNull PyQualifiedExpression element) {
+    // Allow forward references in Pyi annotations
+    if (PyiUtil.isInsideStubAnnotation(element)) {
+      return true;
+    }
+    // Forward references are allowed in annotations according to PEP 563
+    PsiFile file = element.getContainingFile();
+    if (file instanceof PyFile) {
+      final PyFile pyFile = (PyFile)file;
+      return pyFile.getLanguageLevel().isAtLeast(LanguageLevel.PYTHON37) &&
+             pyFile.hasImportFromFuture(FutureFeature.ANNOTATIONS) &&
+             PsiTreeUtil.getParentOfType(element, PyAnnotation.class) != null;
+    }
+    return false;
+  }
 }

python/src/com/jetbrains/python/pyi/PyiReferenceResolveProvider.kt

@@ -1,16 +0,0 @@
-// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
-package com.jetbrains.python.pyi
-
-import com.jetbrains.python.psi.PyQualifiedExpression
-import com.jetbrains.python.psi.resolve.PyReferenceResolveProvider
-import com.jetbrains.python.psi.resolve.RatedResolveResult
-import com.jetbrains.python.psi.types.TypeEvalContext
-
-class PyiReferenceResolveProvider: PyReferenceResolveProvider {
-
-  override fun resolveName(element: PyQualifiedExpression, context: TypeEvalContext): List<RatedResolveResult> = emptyList()
-
-  override fun allowsForwardOutgoingReferencesInClass(element: PyQualifiedExpression): Boolean {
-    return PyiUtil.isInsideStubAnnotation(element)
-  }
-}

python/testData/inspections/PyUnboundLocalVariableInspection/ForwardReferenceInAnnotations.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+def f(a: A): # Should not produce "Name 'A' can be not defined" warning
+    pass
+
+class A:
+    pass

python/testData/pyi/inspections/pyiTopLevelUnboundForwardReferencesInAnnotations/PyiTopLevelUnboundForwardReferencesInAnnotations.pyi

@@ -0,0 +1,7 @@
+def f(x: C) -> C: ... # Should not produce "Name 'C' can be not defined" warning
+
+D = C
+
+class C: ...
+
+E = C

python/testData/resolve/ParameterAnnotationForwardReference.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+class B:
+    def f(self, a: A) -> None:
+        #          <ref>
+        pass
+
+class A:
+    pass

python/testData/resolve/ReturnAnnotationForwardReference.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+class B:
+    def create(self) -> A:
+        #               <ref>
+        pass
+
+class A:
+    pass

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -1264,4 +1264,20 @@ public class PyResolveTest extends PyResolveTestCase {
     assertProjectFilesNotParsed(file);
     assertSdkRootsNotParsed(file);
   }
+
+  // PY-28228
+  public void testReturnAnnotationForwardReference() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () -> assertResolvesTo(PyClass.class, "A")
+    );
+  }
+
+  // PY-28228
+  public void testParameterAnnotationForwardReference() {
+    runWithLanguageLevel(
+      LanguageLevel.PYTHON37,
+      () -> assertResolvesTo(PyClass.class, "A")
+    );
+  }
 }

python/testSrc/com/jetbrains/python/inspections/PyUnboundLocalVariableInspectionTest.java

@@ -201,6 +201,10 @@ public class PyUnboundLocalVariableInspectionTest extends PyInspectionTestCase {
     doTest();
   }
 
+  public void testForwardReferenceInAnnotations() {
+    runWithLanguageLevel(LanguageLevel.PYTHON37, () -> doTest());
+  }
+
   @NotNull
   @Override
   protected Class<? extends PyInspection> getInspectionClass() {

python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java

@@ -117,10 +117,14 @@ public class PyiInspectionsTest extends PyTestCase {
     doPyiTest(PyUnresolvedReferencesInspection.class);
   }
 
-  public void testPyiTopLevelForwardReferencesInAnnotations() {
+  public void testPyiTopLevelUnresolvedForwardReferencesInAnnotations() {
     doPyiTest(PyUnresolvedReferencesInspection.class);
   }
 
+  public void testPyiTopLevelUnboundForwardReferencesInAnnotations() {
+    doPyiTest(PyUnboundLocalVariableInspection.class);
+  }
+
   public void testPyiUnusedImports() {
     doPyiTest(PyUnresolvedReferencesInspection.class);
   }