Fixed bug in extract method refactoring for vars defined before and redefined inside the fragment (PY-6081)

python/src/com/jetbrains/python/codeInsight/codeFragment/PyCodeFragmentBuilder.java

@@ -3,12 +3,15 @@ package com.jetbrains.python.codeInsight.codeFragment;
 import com.intellij.codeInsight.codeFragment.CodeFragmentUtil;
 import com.intellij.codeInsight.codeFragment.Position;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiReference;
 import com.intellij.psi.ResolveResult;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyImportStatementNavigator;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
+import org.jetbrains.annotations.NotNull;
 
 import java.util.*;
 
@@ -136,11 +139,15 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
         }
         // If declaration is before we look for modifications inside
         if (pos == Position.BEFORE) {
+          if (!isTopLevel(element)) {
+            inElements.add(name);
+          }
           final List<PyElement> list = modifiedInsideMap.get(name);
           boolean modificationSeen = false;
           if (list != null) {
             for (PyElement modification : list) {
-              if (modification.getReference().isReferenceTo(declaration)) {
+              final PsiReference reference = modification.getReference();
+              if (reference != null && reference.isReferenceTo(declaration)) {
                 outElements.add(name);
                 modificationSeen = true;
                 break;
@@ -155,6 +162,10 @@ public class PyCodeFragmentBuilder extends PyRecursiveElementVisitor {
     }
   }
 
+  private static boolean isTopLevel(@NotNull PyElement element) {
+    return ScopeUtil.getScopeOwner(element) instanceof PyFile;
+  }
+
   private void processDeclaration(final PyElement element) {
     final Position position = CodeFragmentUtil.getPosition(element, startOffset, endOffset);
     final String name = element.getName();

python/testData/refactoring/extractmethod/LocalVarDefinedBeforeModifiedInside.after.py

@@ -0,0 +1,15 @@
+def bar(foo_new, i_new):
+    need_break = False
+    if i_new > 2:
+        foo_new = False
+        need_break = True
+    return foo_new, need_break
+
+
+def main(indices):
+    foo = True
+    for i in indices:
+        foo, need_break = bar(foo, i)
+        if need_break:
+            break
+    return foo

python/testData/refactoring/extractmethod/LocalVarDefinedBeforeModifiedInside.before.py

@@ -0,0 +1,10 @@
+def main(indices):
+    foo = True
+    for i in indices:
+        <selection>need_break = False
+        if i > 2:
+            foo = False
+            need_break = True</selection>
+        if need_break:
+            break
+    return foo

python/testSrc/com/jetbrains/python/refactoring/PyExtractMethodTest.java

@@ -35,6 +35,11 @@ public class PyExtractMethodTest extends LightMarkedTestCase {
     myFixture.checkResultByFile("/refactoring/extractmethod/" + result);
   }
 
+  private void doTest(String newName) {
+    final String testName = getTestName(false);
+    doTest(testName + ".before.py", newName, testName + ".after.py");
+  }
+
   public void testParameter() {
     doTest("outEmpty/parameter.before.py", "bar", "outEmpty/parameter.after.py");
   }
@@ -152,4 +157,9 @@ public class PyExtractMethodTest extends LightMarkedTestCase {
   public void testMethodInIf() {
     doTest("methodInIf.before.py", "baz", "methodInIf.after.py");
   }
+
+  // PY-6081
+  public void testLocalVarDefinedBeforeModifiedInside() {
+    doTest("bar");
+  }
 }