From 42dd307360ee8fc3231befc87f8647528d06d727 Mon Sep 17 00:00:00 2001
From: Valentina Kiryushkina <valentina.kiryushkina@jetbrains.com>
Date: Fri, 27 Nov 2015 11:18:24 +0300
Subject: [PATCH] Use keyword text range as refrerence text range in order to
 use supreclass handleElementRename. Resolve to  ReferenceExpression in star
 argument of format function

---
 .../PySubstitutionChunkReference.java         | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/PySubstitutionChunkReference.java b/python/src/com/jetbrains/python/codeInsight/PySubstitutionChunkReference.java
index b39f483cbe76..7450f5dd794a 100644
--- a/python/src/com/jetbrains/python/codeInsight/PySubstitutionChunkReference.java
+++ b/python/src/com/jetbrains/python/codeInsight/PySubstitutionChunkReference.java
@@ -16,11 +16,11 @@
 package com.jetbrains.python.codeInsight;
 
 import com.intellij.lang.annotation.HighlightSeverity;
+import com.intellij.openapi.util.TextRange;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiReferenceBase;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtil;
-import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.inspections.PyStringFormatParser;
 import com.jetbrains.python.psi.*;
@@ -32,8 +32,8 @@ public class PySubstitutionChunkReference extends PsiReferenceBase<PyStringLiter
   private final int myPosition;
   private final PyStringFormatParser.SubstitutionChunk myChunk;
 
-  public PySubstitutionChunkReference(PyStringLiteralExpression element, PyStringFormatParser.SubstitutionChunk chunk, final int position) {
-    super(element, chunk.getTextRange().shiftRight(1));
+  public PySubstitutionChunkReference(PyStringLiteralExpression element, @NotNull final PyStringFormatParser.SubstitutionChunk chunk, final int position) {
+    super(element, getKeyWordRange(element, chunk));
     myChunk = chunk;
     myPosition = position;
   }
@@ -49,6 +49,15 @@ public class PySubstitutionChunkReference extends PsiReferenceBase<PyStringLiter
     return null;
   }
 
+  private static TextRange getKeyWordRange(PyStringLiteralExpression element, PyStringFormatParser.SubstitutionChunk chunk) {
+    final TextRange shifted = chunk.getTextRange().shiftRight(1);
+    if (chunk.getMappingKey() != null) {
+      final int start = shifted.getStartOffset() + chunk.getTextRange().substring(element.getStringValue()).indexOf(chunk.getMappingKey());
+      return new TextRange(start, start + chunk.getMappingKey().length());
+    }
+    return shifted;
+  }
+
   @Nullable
   @Override
   public PsiElement resolve() {
@@ -69,8 +78,8 @@ public class PySubstitutionChunkReference extends PsiReferenceBase<PyStringLiter
       PyExpression[] arguments = argumentList.getArguments();
       if (myChunk.getMappingKey() != null) {
         PyKeywordArgument argument = argumentList.getKeywordArgument(myChunk.getMappingKey());
-        PyDictLiteralExpression dictArgument = getDictArgument(arguments);
-        return ObjectUtils.chooseNotNull(argument, dictArgument);
+        PyExpression subStarExpression = getSubStarExpression(arguments);
+        return ObjectUtils.chooseNotNull(argument, subStarExpression);
       }
       else {
         int position = myChunk.getPosition() == null ? myPosition : myChunk.getPosition();
@@ -86,11 +95,6 @@ public class PySubstitutionChunkReference extends PsiReferenceBase<PyStringLiter
     return null;
   }
 
-  @Override
-  public PsiElement handleElementRename(String newElementName) throws IncorrectOperationException {
-    return super.handleElementRename("{" +newElementName + "}");
-  }
-
   private PsiElement resolvePercentString() {
     PyBinaryExpression binaryExpression = PsiTreeUtil.getParentOfType(getElement(), PyBinaryExpression.class);
     final Class[] SIMPLE_RESULT_EXPRESSIONS = {
@@ -136,9 +140,9 @@ public class PySubstitutionChunkReference extends PsiReferenceBase<PyStringLiter
   }
 
   @Nullable
-  private static PyDictLiteralExpression getDictArgument(@NotNull final PyExpression[] args) {
+  private static PyExpression getSubStarExpression(@NotNull final PyExpression[] args) {
     if (args.length == 1 && args[0] instanceof PyStarArgument) {
-      return PsiTreeUtil.getChildOfType(args[0], PyDictLiteralExpression.class);
+      return PsiTreeUtil.getChildOfAnyType(args[0], PyDictLiteralExpression.class, PyReferenceExpression.class);
     }
     return null;
   }