From 4d1b5b0bd0c24da5929dc995fa94c90b8759135b Mon Sep 17 00:00:00 2001
From: Georgii Ustinov <georgii.ustinov@jetbrains.com>
Date: Mon, 18 Dec 2023 15:35:33 +0200
Subject: [PATCH] Migrate AddFixtureToFuncParametersQuickFix and
 PyUnpackTypeVarTupleQuickFix to ModCommand

PY-65297

GitOrigin-RevId: 2695f20eb27475245fb48a93c5ee874041e54463
---
 .../PyUnpackTypeVarTupleQuickFix.java         | 19 ++++++++-----------
 .../PyTestUnpassedFixtureInspection.kt        | 19 ++++++++++---------
 2 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/PyUnpackTypeVarTupleQuickFix.java b/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/PyUnpackTypeVarTupleQuickFix.java
index fd63c45ecd0c..9b3b401961e4 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/PyUnpackTypeVarTupleQuickFix.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/inspections/quickfix/PyUnpackTypeVarTupleQuickFix.java
@@ -1,7 +1,7 @@
 package com.jetbrains.python.inspections.quickfix;
 
-import com.intellij.codeInspection.LocalQuickFix;
-import com.intellij.codeInspection.ProblemDescriptor;
+import com.intellij.modcommand.ModPsiUpdater;
+import com.intellij.modcommand.PsiUpdateModCommandQuickFix;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
@@ -13,7 +13,7 @@ import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PyExpressionStatement;
 import org.jetbrains.annotations.NotNull;
 
-public class PyUnpackTypeVarTupleQuickFix implements LocalQuickFix {
+public class PyUnpackTypeVarTupleQuickFix extends PsiUpdateModCommandQuickFix {
   @Override
   @NotNull
   public String getFamilyName() {
@@ -35,21 +35,18 @@ public class PyUnpackTypeVarTupleQuickFix implements LocalQuickFix {
   }
 
   @Override
-  public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) {
-    PsiElement psiElement = descriptor.getPsiElement();
-    if (psiElement == null) return;
-
-    var languageLevel = LanguageLevel.forElement(psiElement);
+  public void applyFix(@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) {
+    var languageLevel = LanguageLevel.forElement(element);
     PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
 
     if (languageLevel.isAtLeast(LanguageLevel.PYTHON311)) {
-      String starred = "*" + psiElement.getText();
+      String starred = "*" + element.getText();
       PyExpressionStatement expressionStatement = elementGenerator.createFromText(languageLevel, PyExpressionStatement.class, starred);
       PyExpression newElement = expressionStatement.getExpression();
-      psiElement.replace(newElement);
+      element.replace(newElement);
     }
     else {
-      replaceToTypingExtensionsUnpack(psiElement, psiElement, psiElement.getContainingFile(), project);
+      replaceToTypingExtensionsUnpack(element, element, element.getContainingFile(), project);
     }
   }
 }
diff --git a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestUnpassedFixtureInspection.kt b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestUnpassedFixtureInspection.kt
index 3e6b8250cdb6..6dcfb6c638c0 100644
--- a/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestUnpassedFixtureInspection.kt
+++ b/python/src/com/jetbrains/python/testing/pyTestFixtures/PyTestUnpassedFixtureInspection.kt
@@ -1,18 +1,19 @@
 // Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.jetbrains.python.testing.pyTestFixtures
 
-import com.intellij.codeInspection.*
+import com.intellij.codeInspection.LocalInspectionToolSession
+import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.modcommand.ModPsiUpdater
+import com.intellij.modcommand.PsiUpdateModCommandQuickFix
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiElementVisitor
 import com.intellij.psi.util.findParentOfType
 import com.jetbrains.python.PyPsiBundle
 import com.jetbrains.python.inspections.PyInspection
 import com.jetbrains.python.inspections.PyInspectionVisitor
-import com.jetbrains.python.psi.PyDecorator
-import com.jetbrains.python.psi.PyElementGenerator
-import com.jetbrains.python.psi.PyFunction
-import com.jetbrains.python.psi.PyReferenceExpression
-import com.jetbrains.python.psi.PyStringLiteralExpression
+import com.jetbrains.python.psi.*
 import com.jetbrains.python.psi.types.PyFunctionType
 import com.jetbrains.python.psi.types.TypeEvalContext
 
@@ -86,11 +87,11 @@ class PyTestUnpassedFixtureInspection : PyInspection() {
     }
   }
 
-  private inner class AddFixtureToFuncParametersQuickFix : LocalQuickFix {
+  private inner class AddFixtureToFuncParametersQuickFix : PsiUpdateModCommandQuickFix() {
     override fun getFamilyName() = PyPsiBundle.message("QFIX.add.fixture.to.test.function.parameters.list")
 
-    override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-      val element = descriptor.psiElement as? PyReferenceExpression ?: return
+    override fun applyFix(project: Project, element: PsiElement, updater: ModPsiUpdater) {
+      if (element !is PyReferenceExpression) return
       val newParameter = PyElementGenerator.getInstance(project).createParameter(element.text)
       findTestFunc(element)?.parameterList?.addParameter(newParameter)
     }