PY-86756/PY-87353 type representation: support generic scopes and Self types

GitOrigin-RevId: 02686e0ce27872dd22291362f7d578d4e18f4050

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typeRepresentation/psi/PyFunctionTypeRepresentation.kt

@@ -16,23 +16,20 @@
 package com.jetbrains.python.codeInsight.typeRepresentation.psi
 
 import com.intellij.lang.ASTNode
-import com.intellij.psi.PsiElement
 import com.intellij.psi.util.QualifiedName
 import com.jetbrains.python.PyTokenTypes
 import com.jetbrains.python.ast.findChildByClass
 import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
-import com.jetbrains.python.psi.PyClass
 import com.jetbrains.python.psi.PyDoubleStarExpression
 import com.jetbrains.python.psi.PyExpression
-import com.jetbrains.python.psi.PyFile
 import com.jetbrains.python.psi.PyFunction
-import com.jetbrains.python.psi.PyPsiFacade
 import com.jetbrains.python.psi.PyReferenceExpression
 import com.jetbrains.python.psi.PySlashParameter
 import com.jetbrains.python.psi.PyStarExpression
 import com.jetbrains.python.psi.PyTypeParameterList
 import com.jetbrains.python.psi.impl.PyBuiltinCache
 import com.jetbrains.python.psi.impl.PyElementImpl
+import com.jetbrains.python.psi.resolve.PyResolveUtil
 import com.jetbrains.python.psi.types.PyCallableParameter
 import com.jetbrains.python.psi.types.PyCallableParameterImpl
 import com.jetbrains.python.psi.types.PyCallableTypeImpl
@@ -192,34 +189,11 @@ class PyFunctionTypeRepresentation(astNode: ASTNode) : PyElementImpl(astNode), P
     }
   }
 
-  private fun tryResolveFunction(qualifiedFunctionName: QualifiedName, context: TypeEvalContext): PyFunction? {
-    val facade = PyPsiFacade.getInstance(project)
-    val resolveContext = facade.createResolveContextFromFoothold(this)
-
-    // Try to resolve the module first (e.g., "test" from "test.A.f")
-    val moduleName = qualifiedFunctionName.firstComponent?.let { QualifiedName.fromComponents(it) } ?: return null
-    val module = facade.resolveQualifiedName(moduleName, resolveContext).firstOrNull() as? PyFile ?: return null
-
-    // Walk through the remaining components to resolve nested members
-    var current: PsiElement? = module
-    for (i in 1 until qualifiedFunctionName.componentCount) {
-      val componentName = qualifiedFunctionName.components[i] ?: return null
-
-      current = when (val elem = current) {
-        is PyFile -> elem.multiResolveName(componentName).firstOrNull()?.element
-        is PyClass -> {
-          elem.findNestedClass(componentName, false)
-          ?: elem.findMethodByName(componentName, false, context)
-          ?: elem.findClassAttribute(componentName, false, context)
-        }
-        else -> null
-      }
-
-      if (current == null) return null
-    }
-
-    return current as? PyFunction
-  }
+  private fun tryResolveFunction(qualifiedFunctionName: QualifiedName, context: TypeEvalContext): PyFunction? = PyResolveUtil.resolveFullyQualifiedName(
+    qualifiedFunctionName,
+    this,
+    context
+  ) as? PyFunction
 
   private fun resolveTypeExpression(expr: PyExpression, context: TypeEvalContext, typeVarMap: Map<String, PyTypeVarType>): PyType? {
     // Check if this is a reference to a type parameter

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt

@@ -1391,22 +1391,9 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
         if (classType != null) {
           return classType
         }
-        if (context.typeRepresentationMode) {
-          if (resolved.text == "Unknown") {
-            return Ref()
-          }
-          if (resolved is PyFunctionTypeRepresentation) {
-            val result = context.typeContext.getType(resolved)
-            if (result != null) {
-              return Ref(result)
-            }
-          }
-          if (resolved is PySubscriptionExpression) {
-            val moduleType: Ref<PyType?>? = getModuleType(resolved)
-            if (moduleType != null) {
-              return moduleType
-            }
-          }
+        val typeEngineType = getTypeEngineType(typeHint, context)
+        if (typeEngineType != null) {
+          return typeEngineType
         }
         return null
       }
@@ -1420,6 +1407,47 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
       }
     }
 
+    private fun getTypeEngineType(typeHint: PyExpression, context: Context): Ref<PyType?>? {
+      if (!context.typeRepresentationMode) return null
+      if (typeHint.text == "Unknown") {
+        return Ref()
+      }
+      if (typeHint is PyFunctionTypeRepresentation) {
+        val result = context.typeContext.getType(typeHint)
+        if (result != null) {
+          return Ref(result)
+        }
+      }
+      if (typeHint is PySubscriptionExpression) {
+        val moduleType: Ref<PyType?>? = getModuleType(typeHint)
+        if (moduleType != null) {
+          return moduleType
+        }
+      }
+      if (typeHint is PyBinaryExpression && typeHint.operator === PyTokenTypes.AT) {
+        val name: String = typeHint.leftExpression.text
+        val scopeExpression = typeHint.rightExpression!!
+        if (name == SELF) {
+          val type = getType(scopeExpression, context) ?: return null
+
+          val scopeType = type.get()
+          if (scopeType is PyClassType) {
+            return Ref(PySelfType(scopeType))
+          }
+        }
+        val scopeOwner = if (scopeExpression is PyReferenceExpression) {
+          scopeExpression.asQualifiedName()?.let { qualifiedName ->
+            val scopeElement = PyResolveUtil.resolveFullyQualifiedName(qualifiedName, scopeExpression, context.typeContext)
+            scopeElement as? PyQualifiedNameOwner
+          }
+        }
+        else null
+        val result = PyTypeVarTypeImpl(name, null).withScopeOwner(scopeOwner)
+        return Ref(result)
+      }
+      return null
+    }
+
     private fun getModuleType(moduleDefinition: PySubscriptionExpression): Ref<PyType?>? {
       val name = moduleDefinition.rootOperand.name
       if (name == null || name != PyModuleTypeName) {

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

@@ -609,4 +609,75 @@ public final class PyResolveUtil {
       return Collections.emptyList();
     }
   }
+
+  /**
+   * Resolves a fully qualified name like "module.Class.method" by walking through the hierarchy.
+   * This handles nested members (e.g., methods in classes, nested classes) by resolving each component
+   * in sequence starting from the module.
+   *
+   * @param qualifiedName The qualified name to resolve (e.g., "test.A.f")
+   * @param anchor The PSI element to use as context for resolution
+   * @param context The type evaluation context
+   * @return The resolved element, or null if not found
+   */
+  public static @Nullable PsiElement resolveFullyQualifiedName(@NotNull QualifiedName qualifiedName,
+                                                                @NotNull PsiElement anchor,
+                                                                @NotNull TypeEvalContext context) {
+    if (qualifiedName.getComponentCount() == 0) {
+      return null;
+    }
+
+    // Resolve the first component as a module
+    QualifiedName moduleName = QualifiedName.fromComponents(qualifiedName.getFirstComponent());
+    List<PsiElement> moduleResults = PyResolveImportUtil.resolveQualifiedName(
+      moduleName,
+      PyResolveImportUtil.fromFoothold(anchor)
+    );
+
+    PsiElement current = ContainerUtil.getFirstItem(moduleResults);
+    if (current == null) {
+      return null;
+    }
+
+    current = PyUtil.turnDirIntoInit(current);
+
+    // Walk through remaining components to resolve nested members
+    for (String componentName :  qualifiedName.removeHead(1).getComponents()) {
+      if (componentName == null) {
+        return null;
+      }
+
+      if (current instanceof PyFile) {
+        List<RatedResolveResult> members = ((PyFile)current).multiResolveName(componentName);
+        RatedResolveResult firstMember = ContainerUtil.getFirstItem(members);
+        current = firstMember != null ? firstMember.getElement() : null;
+      }
+      else if (current instanceof PyClass) {
+        PyClass pyClass = (PyClass)current;
+        // Try nested class first, then methods
+        PsiElement nestedClass = pyClass.findNestedClass(componentName, false);
+        if (nestedClass != null) {
+          current = nestedClass;
+        }
+        else {
+          PyFunction method = pyClass.findMethodByName(componentName, false, context);
+          if (method != null) {
+            current = method;
+          }
+          else {
+            return null;
+          }
+        }
+      }
+      else {
+        return null;
+      }
+
+      if (current == null) {
+        return null;
+      }
+    }
+
+    return current;
+  }
 }