PY-33055 Fix move attribute to __init__ method quickfix

(cherry picked from commit cd6abb8ad8a01b22066bed85639242b03a85d1d0)

IJ-MR-8050

GitOrigin-RevId: f8a350835ce303c96b7f404263d8367242f4eaa6

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

@@ -398,6 +398,9 @@ QFIX.default.argument=Replace mutable default argument
 #PyMoveAttributeToInitQuickFix
 QFIX.move.attribute=Move attribute to __init__ method
 
+#PyDefineAttributeInInitQuickFix
+QFIX.define.attribute.in.init=Define attribute in __init__ method
+
 #DocstringQuickFix
 QFIX.NAME.docstring=Fix docstring
 QFIX.docstring.add.parameter=Add docstring parameter ''{0}''

python/python-psi-impl/src/com/jetbrains/python/inspections/PyAttributeOutsideInitInspection.java

@@ -17,13 +17,19 @@ package com.jetbrains.python.inspections;
 
 import com.intellij.codeInspection.LocalInspectionToolSession;
 import com.intellij.codeInspection.ProblemsHolder;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiElementVisitor;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
+import com.jetbrains.python.inspections.quickfix.AddFieldQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyMoveAttributeToInitQuickFix;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyClassImpl;
+import com.jetbrains.python.psi.resolve.PyResolveUtil;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.testing.PythonUnitTestDetectorsKt;
 import one.util.streamex.StreamEx;
@@ -101,13 +107,35 @@ public class PyAttributeOutsideInitInspection extends PyInspection {
             !inheritedProperties.contains(attributeName) &&
             !localProperties.containsKey(attributeName) &&
             !isDefinedByProperty(attribute, localProperties.values(), declaredAttributes)) {
-          registerProblem(attribute, PyPsiBundle.message("INSP.attribute.outside.init", attributeName),
-                          new PyMoveAttributeToInitQuickFix());
+          final PyExpression assignedValue = attribute.findAssignedValue();
+          if (assignedValue == null) continue;
+
+          if (expressionReferencesLocalName(assignedValue, node)) {
+            registerProblem(attribute, PyPsiBundle.message("INSP.attribute.outside.init", attributeName),
+                            new AddFieldQuickFix(attributeName, "None", containingClass.getName(), false));
+          }
+          else {
+            registerProblem(attribute, PyPsiBundle.message("INSP.attribute.outside.init", attributeName),
+                            new PyMoveAttributeToInitQuickFix(),
+                            new AddFieldQuickFix(attributeName, "None", containingClass.getName(), false));
+          }
         }
       }
     }
   }
 
+  private static boolean expressionReferencesLocalName(@NotNull PyExpression assignedValue, @NotNull PyFunction function) {
+    Collection<PyReferenceExpression> references = PsiTreeUtil.collectElementsOfType(assignedValue, PyReferenceExpression.class);
+    for (PyReferenceExpression reference : references) {
+      if (reference.isQualified()) continue;
+      Collection<PsiElement> resolved = PyResolveUtil.resolveLocally(reference);
+      if (resolved.isEmpty()) continue;
+      if (ContainerUtil.exists(resolved, it -> it instanceof PyParameter && ((PyParameter)it).isSelf())) continue;
+      if (ContainerUtil.exists(resolved, it -> function == ScopeUtil.getScopeOwner(it))) return true;
+    }
+    return false;
+  }
+
   private static boolean isDefinedByProperty(@NotNull PyTargetExpression attribute,
                                              @NotNull Collection<Property> properties,
                                              @NotNull Map<String, PyTargetExpression> attributesInInit) {

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/addFieldButNotMoveIfImportInMethodBody.py

@@ -0,0 +1,4 @@
+class C:
+    def method(self):
+        import sys
+        self.attr<caret> = sys.path

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/addFieldButNotMoveIfLocalFunctionAssigned.py

@@ -0,0 +1,8 @@
+class Clazz:
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def foo(self):
+        def local_fun():
+            return 42
+        self.x<caret> = 2 + local_fun()

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/addFieldButNotMoveIfLocalVariableAssigned.py

@@ -0,0 +1,7 @@
+class Clazz:
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def foo(self):
+        local_var = 42
+        self.x<caret> = 2 + local_var

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/addFieldButNotMoveIfParameterAssigned.py

@@ -0,0 +1,7 @@
+class Classifier:
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def fit(self, x, y):
+        self.x = x
+        self.y<caret> = y

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/bothAddFieldAndMoveIfAssignedValueHasReferenceToSelf.py

@@ -0,0 +1,7 @@
+class Classifier:
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def fit(self, x):
+        self.x = x
+        self.y<caret> = self.x

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/bothAddFieldAndMoveIfCallAssigned.py

@@ -0,0 +1,10 @@
+class Clazz:
+    def __init__(self, alpha):
+        self.alpha = alpha
+
+    def foo(self):
+        self.x<caret> = self.bar(42)
+
+    def bar(self, x):
+        print(x)
+        return x + 1

python/testData/quickFixes/PyMoveAttributeToInitQuickFixTest/bothAddFieldAndMoveIfTopLevelVarAssigned.py

@@ -0,0 +1,9 @@
+SOME_GLOB_VAL = 42
+
+
+class Clazz:
+    def __init__(self):
+        pass
+
+    def foo(self):
+        self.x<caret> = SOME_GLOB_VAL

python/testSrc/com/jetbrains/python/quickFixes/PyMoveAttributeToInitQuickFixTest.java

@@ -15,11 +15,16 @@
  */
 package com.jetbrains.python.quickFixes;
 
+import com.intellij.codeInsight.intention.IntentionAction;
 import com.intellij.testFramework.TestDataPath;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.PyQuickFixTestCase;
 import com.jetbrains.python.inspections.PyAttributeOutsideInitInspection;
 import com.jetbrains.python.psi.LanguageLevel;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Collection;
+import java.util.List;
 
 @TestDataPath("$CONTENT_ROOT/../testData/quickFixes/PyMoveAttributeToInitQuickFixTest")
 public class PyMoveAttributeToInitQuickFixTest extends PyQuickFixTestCase {
@@ -56,6 +61,74 @@ public class PyMoveAttributeToInitQuickFixTest extends PyQuickFixTestCase {
     });
   }
 
+  // PY-33055
+  public void testAddFieldButNotMoveIfParameterAssigned() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.add.field.to.class", "y", "Classifier")),
+                               List.of(PyPsiBundle.message("QFIX.move.attribute")));
+  }
+
+  // PY-33055
+  public void testBothAddFieldAndMoveIfCallAssigned() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.move.attribute"),
+                                       PyPsiBundle.message("QFIX.add.field.to.class", "x", "Clazz")),
+                               List.of());
+  }
+
+  // PY-33055
+  public void testAddFieldButNotMoveIfLocalVariableAssigned() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.add.field.to.class", "x", "Clazz")),
+                               List.of(PyPsiBundle.message("QFIX.move.attribute")));
+  }
+
+  // PY-33055
+  public void testAddFieldButNotMoveIfLocalFunctionAssigned() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.add.field.to.class", "x", "Clazz")),
+                               List.of(PyPsiBundle.message("QFIX.move.attribute")));
+  }
+
+  // PY-33055
+  public void testBothAddFieldAndMoveIfAssignedValueHasReferenceToSelf() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.move.attribute"),
+                                       PyPsiBundle.message("QFIX.add.field.to.class", "y", "Classifier")),
+                               List.of());
+  }
+
+  // PY-33055
+  public void testBothAddFieldAndMoveIfTopLevelVarAssigned() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.move.attribute"),
+                                       PyPsiBundle.message("QFIX.add.field.to.class", "x", "Clazz")),
+                               List.of());
+  }
+
+  // PY-33055
+  public void testAddFieldButNotMoveIfImportInMethodBody() {
+    doCheckSuggestedQuickFixes(PyAttributeOutsideInitInspection.class,
+                               List.of(PyPsiBundle.message("QFIX.add.field.to.class", "attr", "C")),
+                               List.of(PyPsiBundle.message("QFIX.move.attribute")));
+  }
+
+  protected void doCheckSuggestedQuickFixes(@NotNull Class inspectionClass,
+                                            @NotNull Collection<String> presentHints,
+                                            @NotNull Collection<String> absentHints) {
+    final String testFileName = getTestName(true);
+    myFixture.enableInspections(inspectionClass);
+    myFixture.configureByFile(testFileName + ".py");
+    myFixture.checkHighlighting(true, false, false);
+    for (String hint: presentHints) {
+      myFixture.findSingleIntention(hint);
+    }
+    for (String hint: absentHints) {
+      List<IntentionAction> ints = myFixture.filterAvailableIntentions(hint);
+      assertEmpty(ints);
+    }
+  }
+
   public void testPropertyNegative() {
     doInspectionTest(PyAttributeOutsideInitInspection.class);
   }