PY-54559 PY-80076 Display inlays hints with inferred types for reveal_type and function return types

Merge-request: IJ-MR-157270
Merged-by: Aleksandr Govenko <aleksandr.govenko@jetbrains.com>

GitOrigin-RevId: 1a504fb29b059243ee598672f6f666dbed367b34

python/python-psi-impl/resources/intellij.python.psi.impl.xml

@@ -249,6 +249,17 @@
     <codeInsight.parameterNameHints language="Python"
                                     implementationClass="com.jetbrains.python.inlayHints.PythonInlayParameterHintsProvider"/>
 
+    <codeInsight.declarativeInlayProvider group="TYPES_GROUP"
+                                          bundle="messages.PyPsiBundle"
+                                          implementationClass="com.jetbrains.python.inlayHints.PyTypeInlayHintsProvider"
+                                          isEnabledByDefault="true"
+                                          language="Python"
+                                          nameKey="INLAY.type.hints"
+                                          providerId="python.type.inlays">
+      <option bundle="messages.PyPsiBundle" descriptionKey="INLAY.reveal.type" enabledByDefault="true" nameKey="INLAY.reveal.type" optionId="python.type.inlays.reveal_type" />
+      <option bundle="messages.PyPsiBundle" descriptionKey="INLAY.function.return" enabledByDefault="false" nameKey="INLAY.function.return" optionId="python.type.inlays.function.return" />
+    </codeInsight.declarativeInlayProvider>
+
     <intentionAction>
       <className>com.jetbrains.python.codeInsight.intentions.PyConvertMethodToPropertyIntention</className>
       <bundleName>messages.PyPsiBundle</bundleName>

python/python-psi-impl/resources/messages/PyPsiBundle.properties

@@ -74,6 +74,11 @@ refactoring.inline.can.not.multiline.string.to.f.string=Cannot inline multiline
 
 refactoring.progress.title.updating.existing.usages=Updating existing usages\u2026
 
+### Inlays ###
+INLAY.type.hints=Python Type Hints
+INLAY.reveal.type='reveal_type' argument type
+INLAY.function.return=Function return types
+
 ### Annotators ###
 ANN.deleting.none=Deleting None
 ANN.assign.to.none=Assignment to None

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

@@ -141,6 +141,8 @@ public final class PyTypingTypeProvider extends PyTypeProviderWithCustomContext<
   public static final Pattern TYPE_IGNORE_PATTERN = Pattern.compile("#\\s*type:\\s*ignore\\s*(\\[[^]#]*])?($|(\\s.*))", Pattern.CASE_INSENSITIVE);
 
   public static final String ASSERT_TYPE = "typing.assert_type";
+  public static final String REVEAL_TYPE = "typing.reveal_type";
+  public static final String REVEAL_TYPE_EXT = "typing_extensions.reveal_type";
   public static final String CAST = "typing.cast";
   public static final String CAST_EXT = "typing_extensions.cast";
 

python/resources/inlayProviders/python.type.inlays/preview.py

@@ -0,0 +1,10 @@
+from typing import reveal_type
+
+def example(x: int, y: float)/*<# -> float #>*/:
+    reveal_type(x + y)/*<# float #>*/
+    return x + y
+
+reveal_type(example(1, 2.5))/*<# float #>*/
+    
+    
+    

python/resources/inlayProviders/python.type.inlays/python.type.inlays.function.return.py

@@ -0,0 +1,13 @@
+#  Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+
+from typing import reveal_type
+
+
+def example(x: int, y: float)/*<# -> float #>*/:
+    return x + y
+
+
+def gen(x: list[int])/*<# -> Generator[int, Any, str] #>*/:
+    for i in x:
+        yield i
+    return "end"

python/resources/inlayProviders/python.type.inlays/python.type.inlays.reveal_type.py

@@ -0,0 +1,11 @@
+#  Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+
+from typing import reveal_type
+
+
+def example(x: int, y: float):
+    reveal_type(x + y)/*<# float #>*/
+    return x + y
+
+
+reveal_type(example(1, 2.5))/*<# float #>*/

python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt

@@ -0,0 +1,95 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.inlayHints
+
+import com.intellij.codeInsight.hints.declarative.EndOfLinePosition
+import com.intellij.codeInsight.hints.declarative.HintFontSize
+import com.intellij.codeInsight.hints.declarative.HintFormat
+import com.intellij.codeInsight.hints.declarative.HintMarginPadding
+import com.intellij.codeInsight.hints.declarative.InlayHintsCollector
+import com.intellij.codeInsight.hints.declarative.InlayHintsProvider
+import com.intellij.codeInsight.hints.declarative.InlayTreeSink
+import com.intellij.codeInsight.hints.declarative.InlineInlayPosition
+import com.intellij.codeInsight.hints.declarative.SharedBypassCollector
+import com.intellij.codeInsight.hints.declarative.impl.PresentationTreeBuilderImpl.Companion.MAX_SEGMENT_TEXT_LENGTH
+import com.intellij.openapi.editor.Editor
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REVEAL_TYPE
+import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider.REVEAL_TYPE_EXT
+import com.jetbrains.python.documentation.PythonDocumentationProvider
+import com.jetbrains.python.psi.PyCallExpression
+import com.jetbrains.python.psi.PyFunction
+import com.jetbrains.python.psi.resolve.PyResolveContext
+import com.jetbrains.python.psi.types.PyTypeChecker
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+class PyTypeInlayHintsProvider : InlayHintsProvider {
+  companion object {
+    const val REVEAL_TYPE_OPTION_ID: String = "python.type.inlays.reveal_type"
+    const val FUNCTION_RETURN_TYPE_OPTION_ID: String = "python.type.inlays.function.return"
+  }
+
+  override fun createCollector(file: PsiFile, editor: Editor): InlayHintsCollector? = Collector()
+
+  private class Collector : SharedBypassCollector {
+    val hintFormat = HintFormat.Companion.default
+      .withFontSize(HintFontSize.ABitSmallerThanInEditor)
+      .withHorizontalMargin(HintMarginPadding.MarginAndSmallerPadding)
+
+    override fun collectFromElement(element: PsiElement, sink: InlayTreeSink) {
+      val typeEvalContext = TypeEvalContext.codeAnalysis(element.project, element.containingFile)
+      val resolveContext = PyResolveContext.defaultContext(typeEvalContext)
+
+      sink.whenOptionEnabled(REVEAL_TYPE_OPTION_ID) {
+        getInlaysForRevealType(element, sink, resolveContext)
+      }
+
+      sink.whenOptionEnabled(FUNCTION_RETURN_TYPE_OPTION_ID) {
+        getInlaysForReturnType(element, sink, resolveContext)
+      }
+    }
+
+    private fun getInlaysForRevealType(element: PsiElement, sink: InlayTreeSink, resolveContext: PyResolveContext) {
+      if (element !is PyCallExpression) return
+      val callable = element.multiResolveCalleeFunction(resolveContext).singleOrNull()
+      val typeEvalContext = resolveContext.typeEvalContext
+
+      if (callable is PyFunction && callable.qualifiedName in listOf(REVEAL_TYPE, REVEAL_TYPE_EXT)) {
+        val args = element.getArguments()
+
+        if (args.size != 1) return
+        val type = typeEvalContext.getType(args[0])
+
+        val document = element.containingFile.fileDocument
+        val lineNumber = document.getLineNumber(element.textRange.endOffset)
+        sink.addPresentation(position = EndOfLinePosition(lineNumber), hintFormat = hintFormat) {
+          // use geTypeName here, because reveal_type should show the same as "Type Info" action
+          text(PythonDocumentationProvider.getTypeName(type, typeEvalContext))
+        }
+      }
+    }
+
+    private fun getInlaysForReturnType(element: PsiElement, sink: InlayTreeSink, resolveContext: PyResolveContext) {
+      val typeEvalContext = resolveContext.typeEvalContext
+      val function = element.parent as? PyFunction ?: return
+      if (element == function.nameIdentifier && function.annotationValue == null && function.typeCommentAnnotation == null) {
+        val type = typeEvalContext.getReturnType(function)
+        if (PyTypeChecker.isUnknown(type, typeEvalContext)) return
+        val typeHint = PythonDocumentationProvider.getTypeHint(type, typeEvalContext)
+        sink.addPresentation(position = InlineInlayPosition(function.parameterList.textRange.endOffset, true), hintFormat = hintFormat) {
+          text("-> ")
+          if (typeHint.length >= MAX_SEGMENT_TEXT_LENGTH) {
+            // Platform doesn't allow one text node to be more than 30 characters, but that might not be enough for some types,
+            // for example, 'Generator[str | int, None, int]' is already 31 chars long
+            text(typeHint.substring(0, MAX_SEGMENT_TEXT_LENGTH))
+            text(typeHint.substring(MAX_SEGMENT_TEXT_LENGTH))
+          }
+          else {
+            text(typeHint)
+          }
+        }
+      }
+    }
+  }
+}

python/testSrc/com/jetbrains/python/inlayHints/PyTypeInlayHintsProviderTest.kt

@@ -0,0 +1,64 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.inlayHints
+
+import com.intellij.testFramework.LightProjectDescriptor
+import com.intellij.testFramework.utils.inlays.declarative.DeclarativeInlayHintsProviderTestCase
+import com.jetbrains.python.fixtures.PyLightProjectDescriptor
+import com.jetbrains.python.psi.LanguageLevel
+
+class PyTypeInlayHintsProviderTest : DeclarativeInlayHintsProviderTestCase() {
+  fun testRevealType() {
+    doTest("""
+from typing import reveal_type
+
+def foo(a: int) -> str:
+    reveal_type(a)/*<# int #>*/
+    return "Hi!"
+
+reveal_type(foo(1))/*<# str #>*/
+    """.trimIndent())
+  }
+  
+  fun testFunctionReturnType() {
+    doTest("""
+def foo(a: int) -> str: # no inlay here
+    return "Hi!"
+
+def bar(a: int)/*<# -> str #>*/:
+    return "Hi!"
+    
+def gen(a: int)/*<# -> Generator[int | str | float, Any, str | int] #>*/:
+    yield 42
+    yield "str"
+    yield 42.5
+    if a > 0:
+        return "Hi!"
+    return 42
+    """.trimIndent())
+  }
+  
+  fun testPreview() {
+    doTest("""    
+from typing import reveal_type
+
+def example(x: int, y: float)/*<# -> float #>*/:
+    reveal_type(x + y)/*<# float #>*/
+    return x + y
+
+reveal_type(example(1, 2.5))/*<# float #>*/
+    """.trimIndent())
+  }
+
+  private val allOptions = mapOf(
+    PyTypeInlayHintsProvider.REVEAL_TYPE_OPTION_ID to true,
+    PyTypeInlayHintsProvider.FUNCTION_RETURN_TYPE_OPTION_ID to true,
+  )
+
+  private fun doTest(text: String) {
+    doTestProvider("A.py", text, PyTypeInlayHintsProvider(), allOptions, verifyHintsPresence = true, testMode = ProviderTestMode.SIMPLE)
+  }
+
+  override fun getProjectDescriptor(): LightProjectDescriptor? {
+    return PyLightProjectDescriptor(LanguageLevel.getLatest())
+  }
+}