[python] PY-59986 fix type annotation injection for union operator

Merge-request: IJ-MR-171625
Merged-by: Morgan Bartholomew <morgan.bartholomew@jetbrains.com>

GitOrigin-RevId: 5bd3401e4a98fc615694dacb1d7fe7a1ddb68d3e

python/python-psi-impl/src/com/jetbrains/python/codeInsight/typeHints/PyTypeHintParserDefinition.java

@@ -11,12 +11,6 @@ import com.jetbrains.python.PythonParserDefinition;
 import org.jetbrains.annotations.NotNull;
 
 public final class PyTypeHintParserDefinition extends PythonParserDefinition {
-
-  @Override
-  public @NotNull TokenSet getCommentTokens() {
-    return TokenSet.EMPTY;
-  }
-
   @Override
   public @NotNull PsiFile createFile(@NotNull FileViewProvider viewProvider) {
     return new PyTypeHintFile(viewProvider);

python/src/com/jetbrains/python/codeInsight/PyInjectionUtil.java

@@ -69,8 +69,26 @@ public final class PyInjectionUtil {
   public static @NotNull InjectionResult registerStringLiteralInjection(@NotNull PsiElement element,
                                                                         @NotNull MultiHostRegistrar registrar,
                                                                         @NotNull Language language) {
+    return registerStringLiteralInjection(element, registrar, language, "", "");
+  }
+
+  /**
+   * Registers language injections like {@link #registerStringLiteralInjection} but wraps the entire injected range with parentheses.
+   * This is useful when the injected fragment must be implicitly parenthesized like a type annotation string injection.
+   */
+  public static @NotNull InjectionResult registerStringLiteralInjectionWithParenthesis(@NotNull PsiElement element,
+                                                                        @NotNull MultiHostRegistrar registrar,
+                                                                        @NotNull Language language) {
+    return registerStringLiteralInjection(element, registrar, language, "(", ")");
+  }
+
+  private static @NotNull InjectionResult registerStringLiteralInjection(@NotNull PsiElement element,
+                                                                        @NotNull MultiHostRegistrar registrar,
+                                                                        @NotNull Language language,
+                                                                        @NotNull String prefix,
+                                                                        @NotNull String suffix) {
     registrar.startInjecting(language);
-    final InjectionResult result = processStringLiteral(element, registrar, "", "", Formatting.NONE);
+    final InjectionResult result = processStringLiteral(element, registrar, prefix, suffix, Formatting.NONE);
     if (result.isInjected()) {
       registrar.frankensteinInjection(!result.isStrict())
         .doneInjecting();

python/src/com/jetbrains/python/codeInsight/typing/PyTypingAnnotationInjector.kt

@@ -3,20 +3,16 @@ package com.jetbrains.python.codeInsight.typing
 
 import com.intellij.lang.Language
 import com.intellij.lang.injection.MultiHostRegistrar
-import com.intellij.openapi.util.Condition
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiLanguageInjectionHost
 import com.intellij.psi.util.PsiTreeUtil
-import com.intellij.util.containers.ContainerUtil
 import com.jetbrains.python.codeInsight.PyInjectionUtil
 import com.jetbrains.python.codeInsight.PyInjectorBase
 import com.jetbrains.python.codeInsight.functionTypeComments.PyFunctionTypeAnnotationDialect
 import com.jetbrains.python.codeInsight.typeHints.PyTypeHintDialect
 import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.types.TypeEvalContext
-import com.jetbrains.python.testing.pyTestFixtures.resolve
-import java.util.regex.Pattern
 
 /**
  * Injects fragments for type annotations either in string literals (quoted annotations containing forward references) or
@@ -26,16 +22,25 @@ import java.util.regex.Pattern
 class PyTypingAnnotationInjector : PyInjectorBase() {
   override fun registerInjection(registrar: MultiHostRegistrar, context: PsiElement): PyInjectionUtil.InjectionResult? {
     // Handles only string literals containing quoted types
-    val result = super.registerInjection(registrar, context)
+    getInjectedLanguage(context)?.let { language ->
+      val element = PyInjectionUtil.getLargestStringLiteral(context)
+      if (element != null) {
+        return if (language === PyTypeHintDialect.INSTANCE && "\n" in element.text) {
+          PyInjectionUtil.registerStringLiteralInjectionWithParenthesis(element, registrar, language)
+        }
+        else {
+          PyInjectionUtil.registerStringLiteralInjection(element, registrar, language)
+        }
+      }
+    }
 
-    if (result === PyInjectionUtil.InjectionResult.EMPTY &&
-        context is PsiComment &&
+    if (context is PsiComment &&
         context is PsiLanguageInjectionHost &&
         context.containingFile is PyFile
     ) {
-      return registerCommentInjection(registrar, context as PsiLanguageInjectionHost)
+      return registerCommentInjection(registrar, context)
     }
-    return result
+    return PyInjectionUtil.InjectionResult.EMPTY
   }
 
   override fun getInjectedLanguage(context: PsiElement): Language? {
@@ -72,7 +77,17 @@ class PyTypingAnnotationInjector : PyInjectorBase() {
   }
 
   companion object {
-    val RE_TYPING_ANNOTATION: Pattern = Pattern.compile("\\s*\\S+(\\[.*])?\\s*")
+    val RE_TYPING_ANNOTATION: Regex = Regex(
+      """(?x)
+      \s*
+      \S+(\[.*])?   # initial type like: "list[int]"
+      (\s*\|\s*     # union operator: " | "
+        \S+(\[.*])? # type between union operator
+      )*            # repeating
+      \s*
+      """.trimIndent(),
+      RegexOption.DOT_MATCHES_ALL,
+    )
 
     private fun isInsideValueOfExplicitTypeAnnotation(expr: PyStringLiteralExpression): Boolean {
       val assignment = PsiTreeUtil.getParentOfType(expr, PyAssignmentStatement::class.java)
@@ -148,7 +163,7 @@ class PyTypingAnnotationInjector : PyInjectorBase() {
     }
 
     private fun isTypingAnnotation(s: String): Boolean {
-      return RE_TYPING_ANNOTATION.matcher(s).matches()
+      return RE_TYPING_ANNOTATION matches s
     }
   }
 

python/testData/resolve/QuotedUnionTypeReferenceTopLevel.py

@@ -0,0 +1,2 @@
+a: "Sequence[int] | None"
+               <ref>

python/testSrc/com/jetbrains/python/PyInjectionResolveTest.java

@@ -89,4 +89,9 @@ public class PyInjectionResolveTest extends PyResolveTestCase {
   public void testNewStyleGenericTypeAliasForwardReference() {
     assertResolvesTo(PyClass.class, "MyClass");
   }
+
+  // PY-59986
+  public void testQuotedUnionTypeReferenceTopLevel() {
+    assertResolvesTo(PyClass.class, "int");
+  }
 }