migrate non-typesafe getData: community

GitOrigin-RevId: 9925045d5302a768b5f57d9f965c96c631d14114
2026-03-22 06:50:54 +07:00 · 2024-08-05 16:48:05 +04:00
parent ab92434d9d
commit c98b16d7a2
13 changed files with 90 additions and 69 deletions
--- a/java/java-tests/testSrc/com/intellij/java/psi/search/FindUsagesTargetTest.java
+++ b/java/java-tests/testSrc/com/intellij/java/psi/search/FindUsagesTargetTest.java
@@ -1,9 +1,13 @@
 // Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.java.psi.search;

+import com.intellij.openapi.actionSystem.CommonDataKeys;
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.editor.ex.EditorEx;
 import com.intellij.psi.PsiClass;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiMethod;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase;
@@ -45,7 +49,12 @@ public class FindUsagesTargetTest extends LightJavaCodeInsightFixtureTestCase {
  }

  private PsiElement getTargetElement() {
-    UsageTarget[] targets = UsageTargetUtil.findUsageTargets(((EditorEx)myFixture.getEditor()).getDataContext()::getData);
+    DataContext dataContext = ((EditorEx)myFixture.getEditor()).getDataContext();
+    Editor editor = CommonDataKeys.EDITOR.getData(dataContext);
+    PsiFile psiFile = CommonDataKeys.PSI_FILE.getData(dataContext);
+    PsiElement psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataContext);
+
+    UsageTarget[] targets = UsageTargetUtil.findUsageTargets(editor, psiFile, psiElement);
    assertTrue(targets.length > 0);
    assertTrue(targets[0] instanceof PsiElementUsageTarget);
    return ((PsiElementUsageTarget)targets[0]).getElement();
--- a/platform/core-ui/src/openapi/actionSystem/DataKey.kt
+++ b/platform/core-ui/src/openapi/actionSystem/DataKey.kt
@@ -29,7 +29,7 @@ class DataKey<T> private constructor(override val name: String) : ValueKey<T> {
   */
  fun `is`(dataId: String?): Boolean = name == dataId

-  fun getData(dataContext: DataContext): T? = dataContext.getData(name) as T?
+  fun getData(dataContext: DataContext): T? = dataContext.getData(this)

  fun getData(dataProvider: DataProvider): T? = dataProvider.getData(name) as T?

--- a/platform/lang-impl/src/com/intellij/ide/impl/dataRules/UsageTargetsRule.java
+++ b/platform/lang-impl/src/com/intellij/ide/impl/dataRules/UsageTargetsRule.java
@@ -1,7 +1,11 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.ide.impl.dataRules;

+import com.intellij.openapi.actionSystem.CommonDataKeys;
 import com.intellij.openapi.actionSystem.DataProvider;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.intellij.usages.UsageTargetUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -10,6 +14,9 @@ public final class UsageTargetsRule implements GetDataRule {
  @Override
  @Nullable
  public Object getData(@NotNull DataProvider dataProvider) {
-    return UsageTargetUtil.findUsageTargets(dataProvider);
+    Editor editor = CommonDataKeys.EDITOR.getData(dataProvider);
+    PsiFile file = CommonDataKeys.PSI_FILE.getData(dataProvider);
+    PsiElement psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataProvider);
+    return UsageTargetUtil.findUsageTargets(editor, file, psiElement);
  }
 }
--- a/platform/usageView/api-dump.txt
+++ b/platform/usageView/api-dump.txt
@@ -304,6 +304,7 @@ f:com.intellij.usages.UsageTargetUtil
 - <init>():V
 - s:findUsageTargets(com.intellij.openapi.actionSystem.DataProvider):com.intellij.usages.UsageTarget[]
 - s:findUsageTargets(com.intellij.openapi.editor.Editor,com.intellij.psi.PsiFile):com.intellij.usages.UsageTarget[]
+- s:findUsageTargets(com.intellij.openapi.editor.Editor,com.intellij.psi.PsiFile,com.intellij.psi.PsiElement):com.intellij.usages.UsageTarget[]
 - s:findUsageTargets(com.intellij.psi.PsiElement):com.intellij.usages.UsageTarget[]
 a:com.intellij.usages.UsageToPsiElementProvider
 - <init>():V
--- a/platform/usageView/src/com/intellij/usages/UsageTargetUtil.java
+++ b/platform/usageView/src/com/intellij/usages/UsageTargetUtil.java
@@ -10,6 +10,7 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;

 import java.util.ArrayList;
 import java.util.Collections;
@@ -18,21 +19,27 @@ import java.util.List;
 public final class UsageTargetUtil {
  private static final ExtensionPointName<UsageTargetProvider> EP_NAME = ExtensionPointName.create("com.intellij.usageTargetProvider");

+  /** @deprecated Use {@link #findUsageTargets(Editor, PsiFile, PsiElement)} */
+  @Deprecated(forRemoval = true)
  public static UsageTarget[] findUsageTargets(@NotNull DataProvider dataProvider) {
    Editor editor = CommonDataKeys.EDITOR.getData(dataProvider);
    PsiFile file = CommonDataKeys.PSI_FILE.getData(dataProvider);
+    PsiElement psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataProvider);
+    return findUsageTargets(editor, file, psiElement);
+  }

+  public static UsageTarget @Nullable [] findUsageTargets(@Nullable Editor editor,
+                                                          @Nullable PsiFile file,
+                                                          @Nullable PsiElement psiElement) {
    List<UsageTarget> result = new ArrayList<>();
    if (file != null && editor != null) {
      UsageTarget[] targets = findUsageTargets(editor, file);
      Collections.addAll(result, targets);
    }
-    PsiElement psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataProvider);
    if (psiElement != null) {
      UsageTarget[] targets = findUsageTargets(psiElement);
      Collections.addAll(result, targets);
    }
-
    return result.isEmpty() ? null : result.toArray(UsageTarget.EMPTY_ARRAY);
  }

--- a/platform/xdebugger-impl/src/com/intellij/xdebugger/impl/ui/attach/dialog/AttachToProcessViewWithHosts.kt
+++ b/platform/xdebugger-impl/src/com/intellij/xdebugger/impl/ui/attach/dialog/AttachToProcessViewWithHosts.kt
@@ -33,7 +33,7 @@ abstract class AttachToProcessViewWithHosts(
  companion object{
    // used externally
    @Suppress("MemberVisibilityCanBePrivate")
-    val DEFAULT_ATTACH_HOST: DataKey<String?> = DataKey.create("DEFAULT_ATTACH_HOST")
+    val DEFAULT_ATTACH_HOST: DataKey<String> = DataKey.create("DEFAULT_ATTACH_HOST")
    private fun getDefaultAttachHost(state: AttachDialogState) = state.dataContext.getData(DEFAULT_ATTACH_HOST)
  }

--- a/plugins/gradle/java/src/execution/test/runner/TestTasksChooser.kt
+++ b/plugins/gradle/java/src/execution/test/runner/TestTasksChooser.kt
@@ -2,6 +2,7 @@
 package org.jetbrains.plugins.gradle.execution.test.runner

 import com.intellij.ide.IdeTooltipManager
+import com.intellij.openapi.actionSystem.CustomizedDataContext
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.actionSystem.DataKey
 import com.intellij.openapi.application.ApplicationManager
@@ -150,16 +151,13 @@ open class TestTasksChooser {
    private const val TEST_TASK_NAME = "test"

    @JvmField
-    val LOCATION = DataKey.create<String?>("org.jetbrains.plugins.gradle.execution.test.runner.TestTasksChooser.LOCATION")
+    val LOCATION = DataKey.create<String>("org.jetbrains.plugins.gradle.execution.test.runner.TestTasksChooser.LOCATION")

    @JvmStatic
    fun contextWithLocationName(context: DataContext, locationName: String?): DataContext {
      if (locationName == null) return context
-      return DataContext {
-        when {
-          LOCATION.`is`(it) -> locationName
-          else -> context.getData(it)
-        }
+      return CustomizedDataContext.withSnapshot(context) { sink ->
+        sink[LOCATION] = locationName
      }
    }

--- a/plugins/gradle/java/testSources/importing/GradleImportActionTest.kt
+++ b/plugins/gradle/java/testSources/importing/GradleImportActionTest.kt
@@ -4,6 +4,7 @@ package org.jetbrains.plugins.gradle.importing
 import com.intellij.ide.DataManager
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.CustomizedDataContext
 import org.jetbrains.plugins.gradle.action.ImportProjectFromScriptAction
 import org.junit.Test

@@ -14,13 +15,11 @@ class GradleImportActionTest: GradleImportingTestCase() {

    val action = ImportProjectFromScriptAction()
    val defaultContext = DataManager.getInstance().dataContext
-    val actionEvent = AnActionEvent.createFromDataContext("ProjectViewPopup", null) { dataId ->
-      when {
-        CommonDataKeys.PROJECT.`is`(dataId) -> myProject
-        CommonDataKeys.VIRTUAL_FILE.`is`(dataId) -> virtualFile
-        else -> defaultContext.getData(dataId)
-      }
-    }
+    val actionEvent = AnActionEvent.createFromDataContext(
+      "ProjectViewPopup", null, CustomizedDataContext.withSnapshot(defaultContext) { sink ->
+      sink[CommonDataKeys.PROJECT] = myProject
+      sink[CommonDataKeys.VIRTUAL_FILE] = virtualFile
+    })

    action.update(actionEvent)
    assertFalse(actionEvent.presentation.isVisible)
@@ -34,13 +33,11 @@ class GradleImportActionTest: GradleImportingTestCase() {

    val action = ImportProjectFromScriptAction()
    val defaultContext = DataManager.getInstance().dataContext
-    val actionEvent = AnActionEvent.createFromDataContext("ProjectViewPopup", null) { dataId ->
-      when {
-        CommonDataKeys.PROJECT.`is`(dataId) -> myProject
-        CommonDataKeys.VIRTUAL_FILE.`is`(dataId) -> config
-        else -> defaultContext.getData(dataId)
-      }
-    }
+    val actionEvent = AnActionEvent.createFromDataContext(
+      "ProjectViewPopup", null, CustomizedDataContext.withSnapshot(defaultContext) { sink ->
+      sink[CommonDataKeys.PROJECT] = myProject
+      sink[CommonDataKeys.VIRTUAL_FILE] = config
+    })

    action.update(actionEvent)
    assertTrue("Action should be visible if project is not imported", actionEvent.presentation.isVisible)
--- a/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/findUsages/KotlinFindUsagesTargetTest.kt
+++ b/plugins/kotlin/fir/tests/test/org/jetbrains/kotlin/idea/fir/findUsages/KotlinFindUsagesTargetTest.kt
@@ -1,6 +1,7 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.fir.findUsages

+import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.usages.PsiElementUsageTarget
 import com.intellij.usages.UsageTargetUtil
@@ -25,9 +26,13 @@ class KotlinFindUsagesTargetTest: KotlinLightCodeInsightFixtureTestCase() {

    fun testPrimaryConstructorParam() {
        myFixture.configureByText("Foo.kt", "class Foo(<caret>x: Int)")
-        val targets = UsageTargetUtil.findUsageTargets((myFixture.editor as EditorEx).getDataContext()::getData)
+        val dataContext = (myFixture.editor as EditorEx).getDataContext()
+        val editor = CommonDataKeys.EDITOR.getData(dataContext)
+        val psiFile = CommonDataKeys.PSI_FILE.getData(dataContext)
+        val psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataContext)
+        val targets = UsageTargetUtil.findUsageTargets(editor, psiFile, psiElement)
        assertNotNull(targets)
-        assertEquals(1, targets.size)
+        assertEquals(1, targets!!.size)
        assertInstanceOf(targets[0], PsiElementUsageTarget::class.java)
        val element = (targets[0] as PsiElementUsageTarget).element
        assertTrue(element is KtParameter)
@@ -36,9 +41,13 @@ class KotlinFindUsagesTargetTest: KotlinLightCodeInsightFixtureTestCase() {

    private fun doTestPosition(fileText: String) {
        myFixture.configureByText("Foo.kt", fileText)
-        val targets = UsageTargetUtil.findUsageTargets((myFixture.editor as EditorEx).getDataContext()::getData)
+        val dataContext = (myFixture.editor as EditorEx).getDataContext()
+        val editor = CommonDataKeys.EDITOR.getData(dataContext)
+        val psiFile = CommonDataKeys.PSI_FILE.getData(dataContext)
+        val psiElement = CommonDataKeys.PSI_ELEMENT.getData(dataContext)
+        val targets = UsageTargetUtil.findUsageTargets(editor, psiFile, psiElement)
        assertNotNull(targets)
-        assertEquals(1, targets.size)
+        assertEquals(1, targets!!.size)
        assertInstanceOf(targets[0], PsiElementUsageTarget::class.java)
        val element = (targets[0] as PsiElementUsageTarget).element
        assertTrue(element is KtPrimaryConstructor)
--- a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/refactoring/rename/RenameJvmNameHandler.kt
+++ b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/refactoring/rename/RenameJvmNameHandler.kt
@@ -3,6 +3,7 @@
 package org.jetbrains.kotlin.idea.refactoring.rename

 import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.CustomizedDataContext
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
@@ -53,9 +54,8 @@ class RenameJvmNameHandler : PsiElementRenameHandler() {

                else -> return null
            }
-        return DataContext { id ->
-            if (CommonDataKeys.PSI_ELEMENT.`is`(id)) return@DataContext newElement
-            dataContext.getData(id)
+        return CustomizedDataContext.withSnapshot(dataContext) { sink ->
+            sink.lazy(CommonDataKeys.PSI_ELEMENT) { newElement }
        }
    }

--- a/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/rename/AbstractReferenceSubstitutionRenameHandler.kt
+++ b/plugins/kotlin/refactorings/kotlin.refactorings.common/src/org/jetbrains/kotlin/idea/refactoring/rename/AbstractReferenceSubstitutionRenameHandler.kt
@@ -3,6 +3,7 @@
 package org.jetbrains.kotlin.idea.refactoring.rename

 import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.CustomizedDataContext
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
@@ -49,9 +50,8 @@ abstract class AbstractReferenceSubstitutionRenameHandler(

    override fun invoke(project: Project, editor: Editor?, file: PsiFile?, dataContext: DataContext) {
        val elementToRename = getElementToRename(dataContext) ?: return
-        val wrappingContext = DataContext { id ->
-            if (CommonDataKeys.PSI_ELEMENT.`is`(id)) return@DataContext elementToRename
-            dataContext.getData(id)
+        val wrappingContext = CustomizedDataContext.withSnapshot(dataContext) { sink ->
+            sink.lazy(CommonDataKeys.PSI_ELEMENT) { elementToRename }
        }
        // Can't provide new name for inplace refactoring in unit test mode
        if (!isUnitTestMode() && delegateHandler.isAvailableOnDataContext(wrappingContext)) {
--- a/plugins/maven/src/test/java/org/jetbrains/idea/maven/project/importing/MavenProjectsManagerTest.kt
+++ b/plugins/maven/src/test/java/org/jetbrains/idea/maven/project/importing/MavenProjectsManagerTest.kt
@@ -5,6 +5,7 @@ import com.intellij.ide.DataManager
 import com.intellij.ide.actions.DeleteAction
 import com.intellij.ide.projectView.ProjectView
 import com.intellij.maven.testFramework.MavenMultiVersionImportingTestCase
+import com.intellij.openapi.actionSystem.CustomizedDataContext
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.actionSystem.LangDataKeys
 import com.intellij.openapi.actionSystem.PlatformDataKeys
@@ -725,17 +726,10 @@ class MavenProjectsManagerTest : MavenMultiVersionImportingTestCase() {
  companion object {
    private fun createTestModuleDataContext(vararg modules: Module): DataContext {
      val defaultContext = DataManager.getInstance().getDataContext()
-      return DataContext { dataId: String? ->
-        if (LangDataKeys.MODULE_CONTEXT_ARRAY.`is`(dataId)) {
-          return@DataContext modules
-        }
-        if (ProjectView.UNLOADED_MODULES_CONTEXT_KEY.`is`(dataId)) {
-          return@DataContext listOf<Any>() // UnloadedModuleDescription
-        }
-        if (PlatformDataKeys.DELETE_ELEMENT_PROVIDER.`is`(dataId)) {
-          return@DataContext MavenModuleDeleteProvider()
-        }
-        defaultContext.getData(dataId!!)
+      return CustomizedDataContext.withSnapshot(defaultContext) { sink ->
+        sink[LangDataKeys.MODULE_CONTEXT_ARRAY] = modules
+        sink[ProjectView.UNLOADED_MODULES_CONTEXT_KEY] = listOf() // UnloadedModuleDescription
+        sink[PlatformDataKeys.DELETE_ELEMENT_PROVIDER] = MavenModuleDeleteProvider()
      }
    }
  }
--- a/plugins/maven/testFramework/src/com/intellij/maven/testFramework/MavenDomTestCase.kt
+++ b/plugins/maven/testFramework/src/com/intellij/maven/testFramework/MavenDomTestCase.kt
@@ -14,6 +14,8 @@ import com.intellij.codeInsight.lookup.LookupElementPresentation
 import com.intellij.find.findUsages.PsiElement2UsageTargetAdapter
 import com.intellij.lang.annotation.HighlightSeverity
 import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.actionSystem.CustomizedDataContext
+import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.application.EDT
 import com.intellij.openapi.application.readAction
 import com.intellij.openapi.editor.Editor
@@ -34,7 +36,6 @@ import com.intellij.refactoring.rename.RenameHandler
 import com.intellij.refactoring.rename.RenameHandlerRegistry
 import com.intellij.refactoring.rename.inplace.VariableInplaceRenameHandler
 import com.intellij.refactoring.util.CommonRefactoringUtil.RefactoringErrorHintException
-import com.intellij.testFramework.MapDataContext
 import com.intellij.testFramework.fixtures.CodeInsightTestFixture
 import com.intellij.testFramework.fixtures.CodeInsightTestUtil
 import com.intellij.testFramework.fixtures.IdeaTestFixtureFactory
@@ -491,29 +492,24 @@ abstract class MavenDomTestCase : MavenMultiVersionImportingTestCase() {
    }
  }

-  private suspend fun invokeRename(context: MapDataContext, handler: RenameHandler) {
+  private suspend fun invokeRename(context: DataContext, renameHandler: RenameHandler) {
    withContext(Dispatchers.EDT) {
-      handler.invoke(project, PsiElement.EMPTY_ARRAY, context)
+      renameHandler.invoke(project, PsiElement.EMPTY_ARRAY, context)
    }
  }

-  private suspend fun createDataContext(f: VirtualFile): MapDataContext {
-    val context = MapDataContext()
-
+  private suspend fun createRenameDataContext(f: VirtualFile, value: String?): DataContext {
    val editor = getEditor(f)
-    context.put(CommonDataKeys.EDITOR, editor)
-    context.put(CommonDataKeys.PSI_FILE, getTestPsiFile(f))
-    val targetElement = readAction {
-      TargetElementUtil.findTargetElement(editor, TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED or TargetElementUtil.ELEMENT_NAME_ACCEPTED)
+    val psiFile = getTestPsiFile(f)
+    val context = CustomizedDataContext.withSnapshot(DataContext.EMPTY_CONTEXT) { sink ->
+      sink[CommonDataKeys.EDITOR] = editor
+      sink[PsiElementRenameHandler.DEFAULT_NAME] = value
+      sink.lazy(CommonDataKeys.PSI_FILE) { psiFile }
+      sink.lazy(CommonDataKeys.PSI_ELEMENT) {
+        TargetElementUtil.findTargetElement(
+          editor, TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED or TargetElementUtil.ELEMENT_NAME_ACCEPTED)
+      }
    }
-    context.put(CommonDataKeys.PSI_ELEMENT, targetElement)
-
-    return context
-  }
-
-  private suspend fun createRenameDataContext(f: VirtualFile, value: String?): MapDataContext {
-    val context = createDataContext(f)
-    context.put(PsiElementRenameHandler.DEFAULT_NAME, value)
    return context
  }

@@ -526,11 +522,14 @@ abstract class MavenDomTestCase : MavenMultiVersionImportingTestCase() {
  }

  protected suspend fun search(file: VirtualFile): List<PsiElement> {
-    val context = createDataContext(file)
+    val editor = getEditor(file)
+    val psiFile = getTestPsiFile(file)
    return readAction {
-      val targets = UsageTargetUtil.findUsageTargets { dataId: String? -> context.getData(dataId!!) }
-      val target = (targets[0] as PsiElement2UsageTargetAdapter).element
-      val result: List<PsiReference> = ArrayList(ReferencesSearch.search(target).findAll())
+      val psiElement = TargetElementUtil.findTargetElement(
+        editor, TargetElementUtil.REFERENCED_ELEMENT_ACCEPTED or TargetElementUtil.ELEMENT_NAME_ACCEPTED)
+      val targets = UsageTargetUtil.findUsageTargets(editor, psiFile, psiElement)
+      val target = (targets?.firstOrNull() as? PsiElement2UsageTargetAdapter)?.element ?: return@readAction emptyList()
+      val result = ArrayList(ReferencesSearch.search(target).findAll())
      result.map { it.element }
    }
  }