PY-18816 In PyReferenceExpressionImpl require switch to AST permission only for fragments

that, presumably, contain type hints stored in PSI stubs. Previous
attempt to require passing an explicit instance of type eval context
with the origin PSI file together with the resolve context, instead
of just PyResolveContext#defaultContext() or PyResolveContext#noImplicits(),
caused multiple test failures, since it turned out that there is much
more code relying on the fact that if we were able to get hold of
PyReferenceExpression, we're already in a file that was parsed, and,
thus, we can safely perform reference multiresolve there.

Changing all these existing usages is, first, cumbersome and, second,
doesn't protect us from breakage in other subsystems not covered by
tests. This workaround helps us to avoid unstubbing due to presence of
code fragments needed to parse type annotations in stubs preserving
the legacy behavior for the rest of the code.

Additionally, I've reverted 0aa8039a3a
where I explicitly set type eval context in several places
as described.

python/src/com/jetbrains/python/hierarchy/call/PyStaticCallHierarchyUtil.java

@@ -66,9 +66,8 @@ public class PyStaticCallHierarchyUtil {
       public void visitPyCallExpression(PyCallExpression node) {
         super.visitPyCallExpression(node);
 
-        final TypeEvalContext typeEvalContext = TypeEvalContext.userInitiated(element.getProject(), null);
         StreamEx
-          .of(node.multiResolveCalleeFunction(PyResolveContext.defaultContext().withTypeEvalContext(typeEvalContext)))
+          .of(node.multiResolveCalleeFunction(PyResolveContext.defaultContext()))
           .select(PyFunction.class)
           .forEach(callees::add);
       }

python/src/com/jetbrains/python/highlighting/PyRainbowVisitor.kt

@@ -25,7 +25,6 @@ import com.jetbrains.python.codeInsight.controlflow.ScopeOwner
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.resolve.PyResolveContext
-import com.jetbrains.python.psi.types.TypeEvalContext
 
 class PyRainbowVisitor : RainbowVisitor() {
 
@@ -90,9 +89,7 @@ class PyRainbowVisitor : RainbowVisitor() {
     val parent = targetExpression.parent
     if (parent is PyGlobalStatement) return targetExpression.containingFile
     if (parent is PyNonlocalStatement) {
-      val typeEvalContext = TypeEvalContext.codeAnalysis(targetExpression.project, targetExpression.containingFile)
-      val resolveContext = PyResolveContext.defaultContext().withTypeEvalContext(typeEvalContext)
-      val outerResolved = targetExpression.getReference(resolveContext).resolve()
+      val outerResolved = targetExpression.reference.resolve()
       return if (outerResolved is PyTargetExpression) getTargetContext(outerResolved) else null
     }
 

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

@@ -21,7 +21,6 @@ import com.intellij.psi.PsiElementVisitor;
 import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.psi.PyGlobalStatement;
 import com.jetbrains.python.psi.PyTargetExpression;
-import com.jetbrains.python.psi.resolve.PyResolveContext;
 import org.jetbrains.annotations.Nls;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -58,8 +57,7 @@ public class PyGlobalUndefinedInspection extends PyInspection {
       final PyTargetExpression[] globals = node.getGlobals();
 
       for (PyTargetExpression global : globals) {
-        final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(myTypeEvalContext);
-        if (global.getReference(resolveContext).resolve() == global) {
+        if (global.getReference().resolve() == global) {
           registerProblem(global, PyBundle.message("INSP.NAME.global.$0.undefined", global.getName()));
         }
       }

python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java

@@ -97,7 +97,8 @@ public class PyCallExpressionHelper {
     if (callee instanceof PyReferenceExpression) {
       // dereference
       PyReferenceExpression ref = (PyReferenceExpression)callee;
-      resolveResult = ref.followAssignmentsChain(PyResolveContext.noImplicits());
+      final TypeEvalContext context = TypeEvalContext.codeAnalysis(us.getProject(), us.getContainingFile());
+      resolveResult = ref.followAssignmentsChain(PyResolveContext.noImplicits().withTypeEvalContext(context));
       resolved = resolveResult.getElement();
     }
     else {

python/src/com/jetbrains/python/psi/impl/PyNamedParameterImpl.java

@@ -398,8 +398,7 @@ public class PyNamedParameterImpl extends PyBaseElementImpl<PyNamedParameterStub
         public void visitPyTargetExpression(PyTargetExpression node) {
           if (parameterWasReassigned.get()) return;
 
-          final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
-          if (node.getReference(resolveContext).isReferenceTo(PyNamedParameterImpl.this)) {
+          if (node.getReference().isReferenceTo(PyNamedParameterImpl.this)) {
             parameterWasReassigned.set(true);
           }
           else {

python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java

@@ -22,6 +22,7 @@ import com.intellij.openapi.extensions.ExtensionException;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.util.Ref;
 import com.intellij.psi.*;
+import com.intellij.psi.impl.source.resolve.FileContextUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.containers.ContainerUtil;
@@ -285,7 +286,8 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(context);
     final List<PyType> members = new ArrayList<>();
 
-    if (context.maySwitchToAST(this)) {
+    final PsiFile realFile = FileContextUtil.getContextFile(this);
+    if (!(getContainingFile() instanceof PyExpressionCodeFragment) || (realFile != null && context.maySwitchToAST(realFile))) {
       for (PsiElement target : PyUtil.multiResolveTopPriority(getReference(resolveContext))) {
         if (target == this || target == null) {
           continue;

python/src/com/jetbrains/python/refactoring/classes/membersManager/InstanceFieldsManager.java

@@ -21,8 +21,6 @@ import com.jetbrains.NotNullPredicate;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFunctionBuilder;
-import com.jetbrains.python.psi.resolve.PyResolveContext;
-import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
 
@@ -135,9 +133,7 @@ class InstanceFieldsManager extends FieldsManager {
   private static class FieldsOnly extends NotNullPredicate<PyTargetExpression> {
     @Override
     protected boolean applyNotNull(@NotNull final PyTargetExpression input) {
-      final TypeEvalContext typeEvalContext = TypeEvalContext.userInitiated(input.getProject(), null);
-      final PyResolveContext context = PyResolveContext.noImplicits().withTypeEvalContext(typeEvalContext);
-      return input.getReference(context).resolve() instanceof PyTargetExpression;
+      return input.getReference().resolve() instanceof PyTargetExpression;
     }
   }
 }

python/src/com/jetbrains/python/refactoring/classes/membersManager/PropertiesManager.java

@@ -20,8 +20,6 @@ import com.intellij.psi.PsiReference;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.containers.MultiMap;
 import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.resolve.PyResolveContext;
-import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
@@ -187,14 +185,7 @@ class PropertiesManager extends MembersManager<PyElement> {
 
     @Override
     public void visitPyExpression(final PyExpression node) {
-      final PsiReference reference;
-      if (node instanceof PyReferenceOwner) {
-        final TypeEvalContext context = TypeEvalContext.userInitiated(node.getProject(), null);
-        reference = ((PyReferenceOwner)node).getReference(PyResolveContext.noImplicits().withTypeEvalContext(context));
-      }
-      else {
-        reference = node.getReference();
-      }
+      final PsiReference reference = node.getReference();
       if (reference == null) {
         return;
       }

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

@@ -974,6 +974,17 @@ public class PyTypingTest extends PyTestCase {
            "    expr: Alias = g()");
   }
 
+  // TODO same test for variable type comments
+  // PY-18816
+  public void testLocalTypeAliasInFunctionTypeComment() {
+    doTest("int",
+           "def func():\n" +
+           "    Alias = int\n" +
+           "    def g(x):\n" +
+           "        # type: (Alias) -> None\n" +
+           "        expr = x\n");
+  }
+
   private void doTestNoInjectedText(@NotNull String text) {
     myFixture.configureByText(PythonFileType.INSTANCE, text);
     final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());