From 2894089e5a4f1dc32fc512da391b77d7eb95bd07 Mon Sep 17 00:00:00 2001 From: "Aleksandr.Govenko" Date: Fri, 16 May 2025 17:47:34 +0000 Subject: [PATCH] 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 GitOrigin-RevId: 1a504fb29b059243ee598672f6f666dbed367b34 --- .../resources/intellij.python.psi.impl.xml | 11 +++ .../resources/messages/PyPsiBundle.properties | 5 + .../typing/PyTypingTypeProvider.java | 2 + .../python.type.inlays/preview.py | 10 ++ .../python.type.inlays.function.return.py | 13 +++ .../python.type.inlays.reveal_type.py | 11 +++ .../inlayHints/PyTypeInlayHintsProvider.kt | 95 +++++++++++++++++++ .../PyTypeInlayHintsProviderTest.kt | 64 +++++++++++++ 8 files changed, 211 insertions(+) create mode 100644 python/resources/inlayProviders/python.type.inlays/preview.py create mode 100644 python/resources/inlayProviders/python.type.inlays/python.type.inlays.function.return.py create mode 100644 python/resources/inlayProviders/python.type.inlays/python.type.inlays.reveal_type.py create mode 100644 python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt create mode 100644 python/testSrc/com/jetbrains/python/inlayHints/PyTypeInlayHintsProviderTest.kt diff --git a/python/python-psi-impl/resources/intellij.python.psi.impl.xml b/python/python-psi-impl/resources/intellij.python.psi.impl.xml index b909b7d217cc..13a4676029d2 100644 --- a/python/python-psi-impl/resources/intellij.python.psi.impl.xml +++ b/python/python-psi-impl/resources/intellij.python.psi.impl.xml @@ -249,6 +249,17 @@ + + + com.jetbrains.python.codeInsight.intentions.PyConvertMethodToPropertyIntention messages.PyPsiBundle diff --git a/python/python-psi-impl/resources/messages/PyPsiBundle.properties b/python/python-psi-impl/resources/messages/PyPsiBundle.properties index b101772964a5..b3c4b247ca37 100644 --- a/python/python-psi-impl/resources/messages/PyPsiBundle.properties +++ b/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 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 50afe6ea832e..434bc9de3dbc 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 @@ -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"; diff --git a/python/resources/inlayProviders/python.type.inlays/preview.py b/python/resources/inlayProviders/python.type.inlays/preview.py new file mode 100644 index 000000000000..8ad856e3fe1d --- /dev/null +++ b/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 #>*/ + + + \ No newline at end of file diff --git a/python/resources/inlayProviders/python.type.inlays/python.type.inlays.function.return.py b/python/resources/inlayProviders/python.type.inlays/python.type.inlays.function.return.py new file mode 100644 index 000000000000..ac0d78af471c --- /dev/null +++ b/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" diff --git a/python/resources/inlayProviders/python.type.inlays/python.type.inlays.reveal_type.py b/python/resources/inlayProviders/python.type.inlays/python.type.inlays.reveal_type.py new file mode 100644 index 000000000000..e0c5e11fd398 --- /dev/null +++ b/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 #>*/ diff --git a/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt b/python/src/com/jetbrains/python/inlayHints/PyTypeInlayHintsProvider.kt new file mode 100644 index 000000000000..10ca42db60ce --- /dev/null +++ b/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) + } + } + } + } + } +} \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/inlayHints/PyTypeInlayHintsProviderTest.kt b/python/testSrc/com/jetbrains/python/inlayHints/PyTypeInlayHintsProviderTest.kt new file mode 100644 index 000000000000..e9487e6b5ef6 --- /dev/null +++ b/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()) + } +} \ No newline at end of file