[python] Refactor PyCapturePatternImpl. Introduce PyCaptureContext

(cherry picked from commit 2e3fbf4c7d79e6031c7c087e5c7e7e49046587fd)

IJ-MR-168826

GitOrigin-RevId: b87eda39543460451311fc875d6ae3722d671db0

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

@@ -4,7 +4,7 @@ package com.jetbrains.python.psi;
 import com.jetbrains.python.ast.PyAstCaseClause;
 import org.jetbrains.annotations.Nullable;
 
-public interface PyCaseClause extends PyAstCaseClause, PyStatementPart {
+public interface PyCaseClause extends PyAstCaseClause, PyStatementPart, PyCaptureContext {
   @Override
   default @Nullable PyPattern getPattern() {
     return (PyPattern)PyAstCaseClause.super.getPattern();

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

@@ -4,7 +4,12 @@ package com.jetbrains.python.psi;
 import com.jetbrains.python.ast.PyAstClassPattern;
 import org.jetbrains.annotations.NotNull;
 
-public interface PyClassPattern extends PyAstClassPattern, PyPattern {
+import java.util.Set;
+
+public interface PyClassPattern extends PyAstClassPattern, PyPattern, PyCaptureContext {
+  Set<String> SPECIAL_BUILTINS = Set.of(
+    "bool", "bytearray", "bytes", "dict", "float", "frozenset", "int", "list", "set", "str", "tuple");
+  
   @Override
   default @NotNull PyReferenceExpression getClassNameReference() {
     return (PyReferenceExpression)PyAstClassPattern.super.getClassNameReference();

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

@@ -3,5 +3,5 @@ package com.jetbrains.python.psi;
 
 import com.jetbrains.python.ast.PyAstMappingPattern;
 
-public interface PyMappingPattern extends PyAstMappingPattern, PyPattern {
+public interface PyMappingPattern extends PyAstMappingPattern, PyPattern, PyCaptureContext {
 }

python/python-psi-api/src/com/jetbrains/python/psi/PyPattern.kt

@@ -1,43 +1,77 @@
 // Copyright 2000-2021 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;
+package com.jetbrains.python.psi
 
-import com.jetbrains.python.ast.PyAstPattern;
-import com.jetbrains.python.psi.types.PyType;
-import com.jetbrains.python.psi.types.TypeEvalContext;
-import org.jetbrains.annotations.ApiStatus;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
+import com.intellij.psi.util.findParentOfType
+import com.jetbrains.python.ast.PyAstPattern
+import com.jetbrains.python.psi.types.PyType
+import com.jetbrains.python.psi.types.TypeEvalContext
+import org.jetbrains.annotations.ApiStatus
 
-public interface PyPattern extends PyAstPattern, PyTypedElement {
+interface PyPattern : PyAstPattern, PyTypedElement {
   /**
    * Returns the type that would be captured by this pattern when matching.
-   * <p>
+   *
+   *
    * Unlike other PyTypedElements where getType returns their own type, pattern's getType
    * returns the type that would result from a successful match. For example:
    *
-   * <pre>{@code
+   * ```python
    * class Plant: pass
    * class Animal: pass
    * class Dog(Animal): pass
    *
    * x: Dog | Plant
    * match x:
-   *     case Animal():
-   *         # getType returns Dog here, even though the pattern is Animal()
-   * }</pre>
+   * case Animal():
+   * # getType returns Dog here, even though the pattern is Animal()
+   * ```
    *
-   * @see PyCapturePatternImpl#getCaptureType(PyPattern, TypeEvalContext)
+   * @see PyCapturePatternImpl.getCaptureType
    */
-  @Override
-  @Nullable PyType getType(@NotNull TypeEvalContext context, TypeEvalContext.@NotNull Key key);
+  override fun getType(context: TypeEvalContext, key: TypeEvalContext.Key): PyType?
 
   /**
-   * Decides if the set of values described by a pattern is suitable 
+   * Decides if the set of values described by a pattern is suitable
    * to be subtracted (excluded) from a subject type on the negative edge,
    * or if this pattern is too specific.
    */
   @ApiStatus.Experimental
-  default boolean canExcludePatternType(@NotNull TypeEvalContext context) {
-    return true;
+  fun canExcludePatternType(context: TypeEvalContext): Boolean {
+    return true
+  }
+}
+
+interface PyCaptureContext : PyElement {
+  fun getCaptureTypeForChild(pattern: PyPattern, context: TypeEvalContext): PyType?
+  
+  companion object {
+    /**
+     * Determines what type this pattern would have if it was a capture pattern (like a bare name or _).
+     *
+     * In pattern matching, a capture pattern takes on the type of the entire matched expression,
+     * regardless of any specific pattern constraints.
+     * 
+     * For example:
+     * ```python
+     * x: int | str
+     * match x:
+     *     case a:         # This is a capture pattern
+     *         # Here 'a' has type int | str
+     *     case str():     # This is a class pattern
+     *         # Capture type: int | str (same as what 'case a:' would get)
+     *         # Regular getType: str
+     *
+     * y: int
+     * match y:
+     *     case str() as a:
+     *         # Capture type: int (same as what 'case a:' would get)
+     *         # Regular getType: intersect(int, str) (just 'str' for now)
+     * ```
+     * @see PyPattern#getType(TypeEvalContext, TypeEvalContext.Key)
+     */
+    @JvmStatic
+    fun getCaptureType(pattern: PyPattern, context: TypeEvalContext): PyType? {
+      return pattern.findParentOfType<PyCaptureContext>()?.getCaptureTypeForChild(pattern, context)
+    }
   }
 }

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

@@ -8,7 +8,7 @@ import java.util.List;
 
 import static com.jetbrains.python.ast.PyAstElementKt.findChildrenByClass;
 
-public interface PySequencePattern extends PyAstSequencePattern, PyPattern {
+public interface PySequencePattern extends PyAstSequencePattern, PyPattern, PyCaptureContext {
   default @NotNull List<@NotNull PyPattern> getElements() {
     return List.of(findChildrenByClass(this, PyPattern.class));
   }

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

@@ -2,12 +2,8 @@
 package com.jetbrains.python.psi;
 
 import com.jetbrains.python.ast.PyAstSingleStarPattern;
-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 java.util.List;
 import java.util.Objects;
 
 import static com.jetbrains.python.ast.PyAstElementKt.findChildByClass;
@@ -17,6 +13,4 @@ public interface PySingleStarPattern extends PyAstSingleStarPattern, PyPattern {
   default PyPattern getPattern() {
     return Objects.requireNonNull(findChildByClass(this, PyPattern.class));
   }
-
-  @NotNull List<@Nullable PyType> getCapturedTypesFromSequenceType(@Nullable PyType sequenceType, @NotNull TypeEvalContext context);
 }