PY-24729 Avoid unstubbing while accessing ahead of time annotations

Namely, don't use Scope to find an annotation if AST access is not
allowed.

Added tests on handling of top-level variable annotations and class
attributes annotations in other files. Supposedly, we don't resolve
to unqualified target expressions in other scopes when unstubbing is
forbidden (ordinary local variables are not preserved in PSI stubs).

python/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java

@@ -362,14 +362,26 @@ public class PyTypingTypeProvider extends PyTypeProviderBase {
         }
       }
       else {
-        final Scope scope = ControlFlowCache.getScope(scopeOwner);
-        return StreamEx.of(scope.getNamedElements(name, false))
-          .select(PyTargetExpression.class)
-          .map(x -> getTypeFromTargetExpressionAnnotation(x, context))
-          .nonNull()
-          .map(Ref::get)
-          .findFirst()
-          .orElse(null);
+        StreamEx<PyTargetExpression> candidates = null;
+        if (context.maySwitchToAST(target)) {
+          final Scope scope = ControlFlowCache.getScope(scopeOwner);
+          candidates = StreamEx.of(scope.getNamedElements(name, false)).select(PyTargetExpression.class);
+        }
+        // Unqualified target expression in either class or module
+        else if (scopeOwner instanceof PyFile) {
+          candidates = StreamEx.of(((PyFile)scopeOwner).getTopLevelAttributes()).filter(t -> name.equals(t.getName()));
+        }
+        else if (scopeOwner instanceof PyClass) {
+          candidates = StreamEx.of(((PyClass)scopeOwner).getClassAttributes()).filter(t -> name.equals(t.getName()));
+        }
+        if (candidates != null) {
+          return candidates
+            .map(x -> getTypeFromTargetExpressionAnnotation(x, context))
+            .nonNull()
+            .map(Ref::get)
+            .findFirst()
+            .orElse(null);
+        }
       }
     }
     return null;

python/testData/types/ClassAttributeAnnotationAheadOfTimeInAnotherFile/other.py

+class C:
+    attr: int
+    attr, _ = None, None

python/testData/types/TopLevelVariableAnnotationAheadOfTimeInAnotherFileForTarget/other.py

+x: int
+for x in foo():
+    expr = x

python/testData/types/TopLevelVariableAnnotationAheadOfTimeInAnotherFileUnpackingTarget/other.py

+x: int
+x, y = foo()
+expr = x
\ No newline at end of file

python/testData/types/TopLevelVariableAnnotationAheadOfTimeInAnotherFileWithTarget/other.py

+x: int
+with foo() as x:
+    expr = x

python/testSrc/com/jetbrains/python/PyTypingTest.java

@@ -729,6 +729,13 @@ public class PyTypingTest extends PyTestCase {
   }
 
   // PY-21864
+  public void testTopLevelVariableAnnotationAheadOfTimeInAnotherFileWithTarget() {
+    doMultiFileStubAwareTest("int",
+                             "from other import x\n" +
+                             "\n" +
+                             "expr = x");
+  }
+
   public void testLocalVariableAnnotationAheadOfTimeForTarget() {
     doTest("int",
            "x: int\n" +
@@ -736,6 +743,14 @@ public class PyTypingTest extends PyTestCase {
            "    expr = x\n");
   }
 
+  // PY-21864
+  public void testTopLevelVariableAnnotationAheadOfTimeInAnotherFileForTarget() {
+    doMultiFileStubAwareTest("int",
+                             "from other import x\n" +
+                             "\n" +
+                             "expr = x");
+  }
+
   // PY-21864
   public void testLocalVariableAnnotationAheadOfTimeUnpackingTarget() {
     doTest("int",
@@ -744,6 +759,14 @@ public class PyTypingTest extends PyTestCase {
            "expr = x");
   }
 
+  // PY-21864
+  public void testTopLevelVariableAnnotationAheadOfTimeInAnotherFileUnpackingTarget() {
+    doMultiFileStubAwareTest("int",
+                             "from other import x\n" +
+                             "\n" +
+                             "expr = x");
+  }
+
   // PY-21864
   public void testLocalVariableAnnotationAheadOfTimeOnlyFirstHintConsidered() {
     doTest("int",
@@ -754,6 +777,14 @@ public class PyTypingTest extends PyTestCase {
            "expr = x");
   }
 
+  // PY-21864
+  public void testClassAttributeAnnotationAheadOfTimeInAnotherFile() {
+    doMultiFileStubAwareTest("int",
+                             "from other import C\n" +
+                             "\n" +
+                             "expr = C().attr");
+  }
+
   public void testInstanceAttributeAnnotation() {
     doTest("int",
            "class C:\n" +
@@ -1047,9 +1078,9 @@ public class PyTypingTest extends PyTestCase {
   // PY-24729
   public void testAnnotatedInstanceAttributeInOtherFile() {
     doMultiFileStubAwareTest("int",
-                    "from other import C\n" +
-                    "\n" +
-                    "expr = C().attr");
+                             "from other import C\n" +
+                             "\n" +
+                             "expr = C().attr");
   }
 
   private void doTestNoInjectedText(@NotNull String text) {