fix target element detection for Inline Local (PY-1585); rewrite test so that it doesn't introduce logic which is not there in production

python/src/com/jetbrains/python/refactoring/inline/PyInlineLocalHandler.java

@@ -11,6 +11,7 @@ import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.editor.colors.EditorColors;
 import com.intellij.openapi.editor.colors.EditorColorsManager;
 import com.intellij.openapi.editor.markup.TextAttributes;
+import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.wm.WindowManager;
@@ -44,6 +45,10 @@ public class PyInlineLocalHandler extends InlineActionHandler {
   private static final String REFACTORING_NAME = RefactoringBundle.message("inline.variable.title");
   private static final Pair<PyStatement, Boolean> EMPTY_DEF_RESULT = Pair.create(null, false);
 
+  public static PyInlineLocalHandler getInstance() {
+    return Extensions.findExtension(EP_NAME, PyInlineLocalHandler.class);
+  }
+
   @Override
   public boolean isEnabledForLanguage(Language l) {
     return l instanceof PythonLanguage;
@@ -60,7 +65,13 @@ public class PyInlineLocalHandler extends InlineActionHandler {
       return;
     }
     final PsiReference psiReference = TargetElementUtilBase.findReference(editor);
-    final PyReferenceExpression refExpr = psiReference instanceof PyReferenceExpression ? ((PyReferenceExpression)psiReference) : null;
+    PyReferenceExpression refExpr = null;
+    if (psiReference != null) {
+      final PsiElement refElement = psiReference.getElement();
+      if (refElement instanceof PyReferenceExpression) {
+        refExpr = (PyReferenceExpression) refElement;
+      }
+    }
     invoke(project, editor, (PyTargetExpression)element, refExpr);
   }
 

python/testData/refactoring/inlinelocal/py1585.after.py

@@ -0,0 +1,6 @@
+def foo(a, b):
+    if a < b:
+        params = {'a': 'b'}
+        print params
+    else:
+        print {'c': 'd'}

python/testData/refactoring/inlinelocal/py1585.before.py

@@ -0,0 +1,7 @@
+def foo(a, b):
+    if a < b:
+        params = {'a': 'b'}
+        print params
+    else:
+        params = {'c': 'd'}
+        print p<caret>arams

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

@@ -1,37 +1,27 @@
 package com.jetbrains.python.refactoring;
 
+import com.intellij.codeInsight.TargetElementUtilBase;
 import com.intellij.openapi.util.Comparing;
 import com.intellij.psi.PsiElement;
-import com.jetbrains.python.fixtures.LightMarkedTestCase;
-import com.jetbrains.python.psi.PyReferenceExpression;
-import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.fixtures.PyLightFixtureTestCase;
 import com.jetbrains.python.refactoring.inline.PyInlineLocalHandler;
 
-import java.util.Map;
-
 /**
  * @author Dennis.Ushakov
  */
-public class PyInlineLocalTest extends LightMarkedTestCase {
+public class PyInlineLocalTest extends PyLightFixtureTestCase {
   private void doTest() {
     doTest(null);
   }
 
   private void doTest(String expectedError) {
     final String name = getTestName(true);
-    final Map<String,PsiElement> map = configureByFile("/refactoring/inlinelocal/" + name + ".before.py");
+    myFixture.configureByFile("/refactoring/inlinelocal/" + name + ".before.py");
     try {
-      PsiElement element = map.values().iterator().next().getParent();
-      PyReferenceExpression ref = null;
-      while (element instanceof PyReferenceExpression) {
-        ref = (PyReferenceExpression)element;
-        PsiElement newElement = ((PyReferenceExpression)element).getReference().resolve();
-        if (element == newElement) {
-          break;
-        }
-        element = newElement;
-      }
-      PyInlineLocalHandler.invoke(myFixture.getProject(), myFixture.getEditor(), (PyTargetExpression)element, ref);
+      PsiElement element = TargetElementUtilBase.findTargetElement(myFixture.getEditor(),
+                                                                   TargetElementUtilBase.getInstance().getReferenceSearchFlags());
+      PyInlineLocalHandler handler = PyInlineLocalHandler.getInstance();
+      handler.inlineElement(myFixture.getProject(), myFixture.getEditor(), element);
       if (expectedError != null) fail("expected error: '" + expectedError + "', got none");
     }
     catch (Exception e) {
@@ -67,4 +57,8 @@ public class PyInlineLocalTest extends LightMarkedTestCase {
   public void testPy994() {
     doTest();
   }
+
+  public void testPy1585() {
+    doTest();
+  }
 }