From e044f9e441c1027f0405e52bd1daa71d042347d2 Mon Sep 17 00:00:00 2001
From: Ilia Zakoulov <ilia.zakoulov@jetbrains.com>
Date: Wed, 11 Dec 2024 20:21:05 +0100
Subject: [PATCH] PY-39384: Introduce PyTypeHintProvider

`PyTypeHintProvider` allows providing type hints for resolved elements before the main logic of `PyTypingTypeProvider`.
See an example of `django-stubs`: https://github.com/typeddjango/django-stubs/pull/2335.
`_UserModel` type might be replaced with a custom user model, provided in `settings.py` (`AUTH_USER_MODEL`).

GitOrigin-RevId: 986fb91c800be3ccfbc002c73c673896efec8a1a
---
 .../resources/META-INF/PythonPsi.xml          |  1 +
 .../codeInsight/typing/PyTypeHintProvider.kt  | 25 +++++++++++++++++++
 .../typing/PyTypingTypeProvider.java          |  9 +++++++
 3 files changed, 35 insertions(+)
 create mode 100644 python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt

diff --git a/python/python-psi-api/resources/META-INF/PythonPsi.xml b/python/python-psi-api/resources/META-INF/PythonPsi.xml
index 15585484a3f6..29863960a786 100644
--- a/python/python-psi-api/resources/META-INF/PythonPsi.xml
+++ b/python/python-psi-api/resources/META-INF/PythonPsi.xml
@@ -4,6 +4,7 @@
                     interface="com.jetbrains.python.psi.resolve.PyReferenceResolveProvider"
                     dynamic="true"/>
     <extensionPoint qualifiedName="Pythonid.typeProvider" interface="com.jetbrains.python.psi.impl.PyTypeProvider" dynamic="true"/>
+    <extensionPoint qualifiedName="Pythonid.typeHintProvider" interface="com.jetbrains.python.codeInsight.typing.PyTypeHintProvider" dynamic="true"/>
     <extensionPoint qualifiedName="Pythonid.pySuperMethodsSearch" interface="com.intellij.util.QueryExecutor" dynamic="true"/>
     <extensionPoint qualifiedName="Pythonid.pyClassMembersProvider"
                     interface="com.jetbrains.python.psi.types.PyClassMembersProvider"
diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt
new file mode 100644
index 000000000000..af983cbaf5d2
--- /dev/null
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypeHintProvider.kt
@@ -0,0 +1,25 @@
+package com.jetbrains.python.codeInsight.typing
+
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.util.Ref
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.psi.PyExpression
+import com.jetbrains.python.psi.PyQualifiedNameOwner
+import com.jetbrains.python.psi.types.PyType
+import com.jetbrains.python.psi.types.TypeEvalContext
+import org.jetbrains.annotations.ApiStatus
+import org.jetbrains.annotations.ApiStatus.Experimental
+
+@Experimental
+@ApiStatus.Internal
+interface PyTypeHintProvider {
+  fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref<PyType>?
+
+  companion object {
+    private val EP_NAME: ExtensionPointName<PyTypeHintProvider> = ExtensionPointName.create("Pythonid.typeHintProvider");
+
+    fun parseTypeHint(typeHint: PyExpression, alias: PyQualifiedNameOwner?, resolved: PsiElement, context: TypeEvalContext): Ref<PyType>? {
+      return EP_NAME.extensionList.firstNotNullOfOrNull { it.parseTypeHint(typeHint, alias, resolved, context) }
+    }
+  }
+}
diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
index 4c9931d4b32f..fab73b8ce606 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.java
@@ -865,6 +865,15 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
       context.getTypeAliasStack().add(alias);
     }
     try {
+      final Ref<PyType> typeHintFromProvider = PyTypeHintProvider.Companion.parseTypeHint(
+        typeHint,
+        alias,
+        resolved,
+        context.getTypeContext()
+      );
+      if (typeHintFromProvider != null) {
+        return typeHintFromProvider;
+      }
       final Ref<PyType> typeFromParenthesizedExpression = getTypeFromParenthesizedExpression(resolved, context);
       if (typeFromParenthesizedExpression != null) {
         return typeFromParenthesizedExpression;