From 409ee5968eafd1f89cae165b1c7bfb9bc0b790c3 Mon Sep 17 00:00:00 2001
From: Daniil Kalinin <Daniil.Kalinin@jetbrains.com>
Date: Fri, 17 Nov 2023 13:37:14 +0100
Subject: [PATCH] PY-61858 PEP 695 Type Parameter Syntax: Add Language
 injections inside quoted forward references

GitOrigin-RevId: f38d5d5d3ae9529e2ffb8ddd252f399fb6cea7bd
---
 .../typing/PyTypingAnnotationInjector.java        |  7 +++++++
 .../NewStyleGenericClassBoundForwardReference.py  |  4 ++++
 ...ewStyleGenericFunctionBoundForwardReference.py |  4 ++++
 .../NewStyleGenericTypeAliasForwardReference.py   |  4 ++++
 .../jetbrains/python/PyInjectionResolveTest.java  | 15 +++++++++++++++
 5 files changed, 34 insertions(+)
 create mode 100644 python/testData/resolve/NewStyleGenericClassBoundForwardReference.py
 create mode 100644 python/testData/resolve/NewStyleGenericFunctionBoundForwardReference.py
 create mode 100644 python/testData/resolve/NewStyleGenericTypeAliasForwardReference.py

diff --git a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingAnnotationInjector.java b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingAnnotationInjector.java
index ae3d7802ff38..3deb2d6a5897 100644
--- a/python/src/com/jetbrains/python/codeInsight/typing/PyTypingAnnotationInjector.java
+++ b/python/src/com/jetbrains/python/codeInsight/typing/PyTypingAnnotationInjector.java
@@ -63,6 +63,9 @@ public class PyTypingAnnotationInjector extends PyInjectorBase {
       if (isInsideValueOfExplicitTypeAnnotation(expr)) {
         return PyTypeHintDialect.INSTANCE;
       }
+      if (isInsideNewStyleTypeVarBound(context)) {
+        return PyTypeHintDialect.INSTANCE;
+      }
     }
     return null;
   }
@@ -133,6 +136,10 @@ public class PyTypingAnnotationInjector extends PyInjectorBase {
     return function != null && function.getTypeComment() == comment;
   }
 
+  private static boolean isInsideNewStyleTypeVarBound(@NotNull PsiElement element) {
+    return PsiTreeUtil.getParentOfType(element, PyTypeParameter.class, true, PyTypeParameterListOwner.class) != null;
+  }
+
   private static boolean isTypingAnnotation(@NotNull String s) {
     return RE_TYPING_ANNOTATION.matcher(s).matches();
   }
diff --git a/python/testData/resolve/NewStyleGenericClassBoundForwardReference.py b/python/testData/resolve/NewStyleGenericClassBoundForwardReference.py
new file mode 100644
index 000000000000..4c55946a6ffc
--- /dev/null
+++ b/python/testData/resolve/NewStyleGenericClassBoundForwardReference.py
@@ -0,0 +1,4 @@
+class A[T: "MyClass"]: ...
+            <ref>
+
+class MyClass: ...
\ No newline at end of file
diff --git a/python/testData/resolve/NewStyleGenericFunctionBoundForwardReference.py b/python/testData/resolve/NewStyleGenericFunctionBoundForwardReference.py
new file mode 100644
index 000000000000..472826c1c891
--- /dev/null
+++ b/python/testData/resolve/NewStyleGenericFunctionBoundForwardReference.py
@@ -0,0 +1,4 @@
+def foo[T: "MyClass"](x: T) -> T: ...
+            <ref>
+
+class MyClass: ...
\ No newline at end of file
diff --git a/python/testData/resolve/NewStyleGenericTypeAliasForwardReference.py b/python/testData/resolve/NewStyleGenericTypeAliasForwardReference.py
new file mode 100644
index 000000000000..a87b558f46b5
--- /dev/null
+++ b/python/testData/resolve/NewStyleGenericTypeAliasForwardReference.py
@@ -0,0 +1,4 @@
+type myType[T: "MyClass"] = list[T]
+                <ref>
+
+class MyClass: ...
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyInjectionResolveTest.java b/python/testSrc/com/jetbrains/python/PyInjectionResolveTest.java
index 9ede52a19748..5b2012fe1f13 100644
--- a/python/testSrc/com/jetbrains/python/PyInjectionResolveTest.java
+++ b/python/testSrc/com/jetbrains/python/PyInjectionResolveTest.java
@@ -74,4 +74,19 @@ public class PyInjectionResolveTest extends PyResolveTestCase {
   public void testFunctionTypeCommentReturnTypeReference() {
     assertResolvesTo(PyClass.class, "MyClass");
   }
+
+  // PY-61858
+  public void testNewStyleGenericClassBoundForwardReference() {
+    assertResolvesTo(PyClass.class, "MyClass");
+  }
+
+  // PY-61858
+  public void testNewStyleGenericFunctionBoundForwardReference() {
+    assertResolvesTo(PyClass.class, "MyClass");
+  }
+
+  // PY-61858
+  public void testNewStyleGenericTypeAliasForwardReference() {
+    assertResolvesTo(PyClass.class, "MyClass");
+  }
 }