PY-63367 PY-63366 Include instructions for annotations of new-styles generic functions in their own CFG

This way they have access to type parameters of these functions, and PyReferenceImpl.getResultsFromProcessor
can properly perform flow-sensitive resolve for them now that READ instructions for expressions
inside annotations and WRITE instructions for type parameters are in the same CFG. It allowed to
remove all the special logic for type parameters from PyReferenceImpl as well as from
PyUnboundLocalVariableInspection.

I had to special-case these annotations in PyResolveUtil.scopeCrawlUp, though, to "lift" their
original scope back from what ScopeUtil.getScopeOwner returns, making names defined in a class scope,
e.g., nested classes, visible to them.

GitOrigin-RevId: b2e78211565300dbcd7c2e2b246a7a8e14bb0e8f

python/python-common-tests/com/jetbrains/python/PyCommonResolveTest.java

@@ -1845,6 +1845,21 @@ public abstract class PyCommonResolveTest extends PyCommonResolveTestCase {
     });
   }
 
+  // PY-61877 PY-63366
+  public void testNestedClassNotResolvedInsideMethodOfNewStyleGenericClass() {
+    assertNotResolved();
+  }
+
+  // PY-61877 PY-63367
+  public void testNestedClassResolvedInsideAnnotationOfNewStyleGenericMethod() {
+    assertResolvesTo(PyClass.class, "Nested");
+  }
+
+  // PY-61877
+  public void testNestedClassIsResolvedInSuperclassListOfAnotherNewStyleGenericNestedClass() {
+    assertResolvesTo(PyClass.class, "Nested");
+  }
+
   // [TODO] daniil.kalinin enable when resolve for collisions in type parameter names and class attribute names is implemented
   // PY-61877
   //public void testClassAttributeDeclarationWithSameAsTypeParameterNameNotResolvedToTypeParameter() {

python/python-psi-impl/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -91,9 +91,11 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
     visitParameterListExpressions(node.getParameterList());
     visitDecorators(node.getDecoratorList());
-    final PyAnnotation annotation = node.getAnnotation();
-    if (annotation != null) {
-      annotation.acceptChildren(this);
+    if (node.getTypeParameterList() == null) {
+      final PyAnnotation annotation = node.getAnnotation();
+      if (annotation != null) {
+        annotation.acceptChildren(this);
+      }
     }
 
     final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
@@ -121,9 +123,11 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
         if (defaultValue != null) {
           defaultValue.accept(PyControlFlowBuilder.this);
         }
-        final PyAnnotation annotation = param.getAnnotation();
-        if (annotation != null) {
-          annotation.acceptChildren(PyControlFlowBuilder.this);
+        if (parameterList.getParent() instanceof PyFunction function && function.getTypeParameterList() == null) {
+          final PyAnnotation annotation = param.getAnnotation();
+          if (annotation != null) {
+            annotation.acceptChildren(PyControlFlowBuilder.this);
+          }
         }
       }
     });
@@ -302,6 +306,10 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
   @Override
   public void visitPyNamedParameter(final @NotNull PyNamedParameter node) {
+    PyAnnotation annotation = node.getAnnotation();
+    if (annotation != null) {
+      annotation.accept(this);
+    }
     final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
     myBuilder.addNode(instruction);
     myBuilder.checkPending(instruction);
@@ -309,8 +317,12 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
 
   @Override
   public void visitPyAnnotation(@NotNull PyAnnotation node) {
-    // Return type and parameter annotations for functions are evaluated in their enclosing scope and processed in visitPyFunction
-    if (!(node.getParent() instanceof PyFunction || node.getParent() instanceof PyParameter)) {
+    // Unless there is a type parameter list, return type and parameter annotations for functions are evaluated in their enclosing scope
+    // and processed in visitPyFunction.
+    // If there are type parameters, though, we need to put the corresponding instructions *inside* the function's scope to be able to
+    // access them from annotations. 
+    PyFunction function = PsiTreeUtil.getParentOfType(node, PyFunction.class, true, PyStatement.class);
+    if (function == null || function.getTypeParameterList() != null) {
       super.visitPyAnnotation(node);
     }
   }
@@ -1036,11 +1048,6 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
   public void visitPyTypeAliasStatement(@NotNull PyTypeAliasStatement node) {
     myBuilder.startNode(node);
 
-    PyExpression typeExpression = node.getTypeExpression();
-    if (typeExpression != null) {
-      typeExpression.accept(this);
-    }
-
     final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
     myBuilder.addNode(instruction);
     myBuilder.checkPending(instruction);

python/python-psi-impl/src/com/jetbrains/python/inspections/PyUnboundLocalVariableInspection.java

@@ -126,9 +126,6 @@ public class PyUnboundLocalVariableInspection extends PyInspection {
         if (PyInspectionsUtil.hasAnyInterruptedControlFlowPaths(node)) {
           return;
         }
-        else if (resolved instanceof PyTypeParameter && currentScopeOwner instanceof PyFunction) {
-          return;
-        }
         if (owner instanceof PyFile) {
           if (isBuiltin) {
             return;

python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTypeAliasStatementImpl.java

@@ -14,6 +14,7 @@ import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
 
 public class PyTypeAliasStatementImpl extends PyBaseElementImpl<PyTypeAliasStatementStub> implements PyTypeAliasStatement {
   public PyTypeAliasStatementImpl(ASTNode astNode) {
@@ -109,4 +110,10 @@ public class PyTypeAliasStatementImpl extends PyBaseElementImpl<PyTypeAliasState
     @Nullable PsiElement name = getNameIdentifier();
     return name != null ? name.getTextOffset() : super.getTextOffset();
   }
+
+  @Override
+  public void subtreeChanged() {
+    super.subtreeChanged();
+    ControlFlowCache.clear(this);
+  }
 }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java

@@ -241,11 +241,6 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     final ScopeOwner resolvedOwner = processor.getOwner();
 
     final Collection<PsiElement> resolvedElements = processor.getElements();
-
-    if (referenceOwner instanceof PyTypeParameterListOwner typeParameterListOwner) {
-      tryMatchTypeParametersInDeclaration(referencedName, typeParameterListOwner, realContext, resultList, typeEvalContext);
-    }
-
     if (resolvedOwner != null && !resolvedElements.isEmpty() && !ControlFlowCache.getScope(resolvedOwner).isGlobal(referencedName)) {
       if (resolvedOwner == referenceOwner) {
         final List<Instruction> instructions = getLatestDefinitions(referencedName, resolvedOwner, realContext);
@@ -274,9 +269,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
           resolveInParentScope = () -> ScopeUtil.getScopeOwner(resolvedOwner);
         }
         else {
-          if (resultList.isEmpty()) {
-            unreachableLocalDeclaration = true;
-          }
+          unreachableLocalDeclaration = true;
         }
       }
       else if (referenceOwner != null) {
@@ -317,18 +310,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
             continue;
           }
           if (definer == null) {
-            if (resolved instanceof PyTypeParameter) {
-              if (resolvedOwner != referenceOwner ||
-                  PsiTreeUtil.getParentOfType(realContext, PyAnnotation.class, PyTypeCommentOwner.class) != null) {
-                resultList.poke(resolved, RatedResolveResult.RATE_HIGH);
-              }
-              else {
-                resultList.poke(resolved, RatedResolveResult.RATE_LOW);
-              }
-            }
-            else {
-              resultList.poke(resolved, getRate(resolved, typeEvalContext));
-            }
+            resultList.poke(resolved, getRate(resolved, typeEvalContext));
           }
           else {
             resultList.poke(definer, getRate(definer, typeEvalContext));
@@ -391,22 +373,6 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
     return results;
   }
 
-  private static void tryMatchTypeParametersInDeclaration(@NotNull String referencedName,
-                                                          PyTypeParameterListOwner owner,
-                                                          PsiElement realContext,
-                                                          ResolveResultList resultList,
-                                                          TypeEvalContext typeEvalContext) {
-    PyTypeParameterList typeParameterList = owner.getTypeParameterList();
-    if (typeParameterList != null && PsiTreeUtil.getParentOfType(realContext, PyTypeParameterListOwner.class) != null) {
-      typeParameterList.getTypeParameters()
-        .stream().filter(typeParameter -> typeParameter.getName() != null && typeParameter.getName().equals(referencedName))
-        .forEach(typeParameter -> {
-                   resultList.add(new RatedResolveResult(getRate(typeParameter, typeEvalContext), typeParameter));
-                 }
-        );
-    }
-  }
-
   @NotNull
   private ResolveResultList resolveByReferenceResolveProviders() {
     final ResolveResultList results = new ResolveResultList();

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

@@ -59,8 +59,8 @@ public final class PyResolveUtil {
    */
   public static void scopeCrawlUp(@NotNull PsiScopeProcessor processor, @NotNull PsiElement element, @Nullable String name,
                                   @Nullable PsiElement roof) {
-    final ScopeOwner originalOwner = ScopeUtil.getScopeOwner(element);
-    final PsiElement parent = element.getParent();
+    ScopeOwner originalOwner = ScopeUtil.getScopeOwner(element);
+    PsiElement parent = element.getParent();
     ScopeOwner owner = originalOwner;
     if (parent instanceof PyNonlocalStatement) {
       /* wee need to search in one step out scope for nonlocal statements */
@@ -76,6 +76,24 @@ public final class PyResolveUtil {
         owner = (PyFile)globalScope;
       }
     }
+    if (owner instanceof PyFunction function && function.getTypeParameterList() != null) {
+      // Type parameters of generic functions and methods need to be visible in their parameter and return type annotations,
+      // so we resolve names in these annotations starting from the containing function's scope, "lowering" their scope in
+      // ScopeUtil.getScopeOwner().
+      // At the same time, these annotations still need to see names from the class scope, such as nested classes defined there,
+      // as if these names were evaluated on a class-level scope.
+      PyAnnotation annotation = PsiTreeUtil.getParentOfType(element, PyAnnotation.class);
+      if (PsiTreeUtil.getParentOfType(annotation, PyFunction.class, true, PyStatement.class) == function) {
+        originalOwner = ScopeUtil.getScopeOwner(function);
+      }
+    }
+    if (owner instanceof PyClass pyClass && pyClass.getTypeParameterList() != null) {
+      // The same logic applies to the list of base classes of a generic class
+      PyArgumentList superclassList = PsiTreeUtil.getParentOfType(element, PyArgumentList.class);
+      if (superclassList != null && superclassList.getParent() == pyClass) {
+        originalOwner = ScopeUtil.getScopeOwner(pyClass);
+      }
+    }
     scopeCrawlUp(processor, owner, originalOwner, name, roof);
   }
 
@@ -87,29 +105,30 @@ public final class PyResolveUtil {
   public static void scopeCrawlUp(@NotNull PsiScopeProcessor processor, @Nullable ScopeOwner scopeOwner,
                                   @Nullable ScopeOwner originalScopeOwner, @Nullable String name, @Nullable PsiElement roof) {
     while (scopeOwner != null) {
-      if (!(scopeOwner instanceof PyClass) || scopeOwner == originalScopeOwner || ((PyClass)scopeOwner).getTypeParameterList() != null) {
-        final Scope scope = ControlFlowCache.getScope(scopeOwner);
-        if (name != null) {
-          final boolean includeNestedGlobals = scopeOwner instanceof PyFile;
-          for (PsiNamedElement resolved : scope.getNamedElements(name, includeNestedGlobals)) {
-            if (!processor.execute(resolved, ResolveState.initial())) {
-              return;
-            }
-          }
-        }
-        else {
-          for (PsiNamedElement element : scope.getNamedElements()) {
-            if (!processor.execute(element, ResolveState.initial())) {
-              return;
-            }
-          }
-        }
-        for (PyImportedNameDefiner definer : scope.getImportedNameDefiners()) {
-          if (!processor.execute(definer, ResolveState.initial())) {
+      final Scope scope = ControlFlowCache.getScope(scopeOwner);
+      if (name != null) {
+        final boolean includeNestedGlobals = scopeOwner instanceof PyFile;
+        for (PsiNamedElement resolved : scope.getNamedElements(name, includeNestedGlobals)) {
+          if (isClassLevelDefinitionInvisibleToReference(resolved, scopeOwner, originalScopeOwner)) continue;
+          if (!processor.execute(resolved, ResolveState.initial())) {
             return;
           }
         }
       }
+      else {
+        for (PsiNamedElement element : scope.getNamedElements()) {
+          if (isClassLevelDefinitionInvisibleToReference(element, scopeOwner, originalScopeOwner)) continue;
+          if (!processor.execute(element, ResolveState.initial())) {
+            return;
+          }
+        }
+      }
+      for (PyImportedNameDefiner definer : scope.getImportedNameDefiners()) {
+        if (isClassLevelDefinitionInvisibleToReference(definer, scopeOwner, originalScopeOwner)) continue;
+        if (!processor.execute(definer, ResolveState.initial())) {
+          return;
+        }
+      }
       if (scopeOwner == roof) {
         return;
       }
@@ -122,6 +141,12 @@ public final class PyResolveUtil {
     }
   }
 
+  private static boolean isClassLevelDefinitionInvisibleToReference(@NotNull PsiElement definition,
+                                                                    @NotNull ScopeOwner definitionScope,
+                                                                    @Nullable ScopeOwner referenceScope) {
+    return definitionScope instanceof PyClass pyClass && pyClass != referenceScope && !(definition instanceof PyTypeParameter);
+  }
+
   /**
    * Resolves the passed expression in its containing file.
    * Does not go outside this file.

python/testData/codeInsight/controlflow/NewStyleGenericFunctionAnnotationsAreIncludedInItsGraph.py

@@ -0,0 +1,2 @@
+def f[T](x: list[T] = list()) -> T:
+    return x[0]

python/testData/codeInsight/controlflow/NewStyleGenericFunctionAnnotationsAreIncludedInItsGraph.txt

@@ -0,0 +1,12 @@
+0(1) element: null
+1(2) element: PyTypeParameter
+2(3) WRITE ACCESS: T
+3(4) element: PySubscriptionExpression
+4(5) READ ACCESS: list
+5(6) READ ACCESS: T
+6(7) WRITE ACCESS: x
+7(8) READ ACCESS: T
+8(9) element: PyReturnStatement
+9(10) element: PySubscriptionExpression
+10(11) READ ACCESS: x
+11() element: null

python/testData/codeInsight/controlflow/NewStyleGenericFunctionAnnotationsAreNotIncludedInEnclosingScopeGraph.py

@@ -0,0 +1,2 @@
+def f[T](x: list[T] = list()) -> T:
+    return x[0]

python/testData/codeInsight/controlflow/NewStyleGenericFunctionAnnotationsAreNotIncludedInEnclosingScopeGraph.txt

@@ -0,0 +1,5 @@
+0(1) element: null
+1(2) element: PyFunction('f')
+2(3) READ ACCESS: list
+3(4) WRITE ACCESS: f
+4() element: null

python/testData/codeInsight/controlflow/TypeAliasStatement.txt

@@ -1,5 +1,4 @@
 0(1) element: null
 1(2) element: PyTypeAliasStatement
-2(3) READ ACCESS: str
-3(4) WRITE ACCESS: myType
-4() element: null
+2(3) WRITE ACCESS: myType
+3() element: null

python/testData/codeInsight/controlflow/TypeParameterListInFunctionDeclaration.txt

@@ -3,6 +3,8 @@
 2(3) WRITE ACCESS: T
 3(4) element: PyTypeParameter
 4(5) WRITE ACCESS: U
-5(6) WRITE ACCESS: a
-6(7) WRITE ACCESS: b
-7() element: null
+5(6) READ ACCESS: T
+6(7) WRITE ACCESS: a
+7(8) READ ACCESS: U
+8(9) WRITE ACCESS: b
+9() element: null

python/testData/resolve/NestedClassIsResolvedInSuperclassListOfAnotherNewStyleGenericNestedClass.py

@@ -0,0 +1,7 @@
+class C:
+    class Nested:
+        pass
+        
+    class Another[T](Nested):
+        #                <ref>
+        pass

python/testData/resolve/NestedClassNotResolvedInsideMethodOfNewStyleGenericClass.py

@@ -0,0 +1,7 @@
+class C[T]:
+    class Nested:
+        pass
+        
+    def method():
+        Nested
+#        <ref> 

python/testData/resolve/NestedClassResolvedInsideAnnotationOfNewStyleGenericMethod.py

@@ -0,0 +1,7 @@
+class C:
+    class Nested:
+        pass
+
+    def method[T](self, x: Nested):
+        #                   <ref>     
+        pass

python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java

@@ -549,6 +549,16 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  // PY-61877
+  public void testNewStyleGenericFunctionAnnotationsAreIncludedInItsGraph() {
+    doTestFirstStatement();
+  }
+
+  // PY-61877
+  public void testNewStyleGenericFunctionAnnotationsAreNotIncludedInEnclosingScopeGraph() {
+    doTest();
+  }
+
   private void doTestFirstStatement() {
     final String testName = getTestName(false);
     configureByFile(testName + ".py");