python/src/com/jetbrains/python/refactoring/move/moduleMembers/PyMoveSymbolProcessor.java

   1 /*
   2  * Copyright 2000-2015 JetBrains s.r.o.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  * http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 package com.jetbrains.python.refactoring.move.moduleMembers;
  17
  18 import com.intellij.openapi.util.Condition;
  19 import com.intellij.psi.*;
  20 import com.intellij.psi.util.PsiTreeUtil;
  21 import com.intellij.psi.util.PsiUtilCore;
  22 import com.intellij.psi.util.QualifiedName;
  23 import com.intellij.usageView.UsageInfo;
  24 import com.intellij.util.containers.ContainerUtil;
  25 import com.jetbrains.python.PyNames;
  26 import com.jetbrains.python.codeInsight.PyDunderAllReference;
  27 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
  28 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
  29 import com.jetbrains.python.psi.*;
  30 import com.jetbrains.python.psi.resolve.PyResolveContext;
  31 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
  32 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
  33 import org.jetbrains.annotations.NotNull;
  34 import org.jetbrains.annotations.Nullable;
  35
  36 import java.util.*;
  37
  38 import static com.jetbrains.python.psi.impl.PyImportStatementNavigator.getImportStatementByElement;
  39
  40 /**
  41  * @author Mikhail Golubev
  42  */
  43 public class PyMoveSymbolProcessor {
  44   private final PsiNamedElement myMovedElement;
  45   private final PyFile myDestinationFile;
  46   private final List<UsageInfo> myUsages;
  47   private final PsiElement[] myAllMovedElements;
  48   private final List<PsiFile> myOptimizeImportTargets = new ArrayList<>();
  49   private final Set<ScopeOwner> myScopeOwnersWithGlobal = new HashSet<>();
  50
  51   public PyMoveSymbolProcessor(@NotNull final PsiNamedElement element,
  52                                @NotNull PyFile destination,
  53                                @NotNull Collection<UsageInfo> usages,
  54                                @NotNull PsiElement[] otherElements) {
  55     myMovedElement = element;
  56     myDestinationFile = destination;
  57     myAllMovedElements = otherElements;
  58     myUsages = ContainerUtil.sorted(usages, (u1, u2) -> PsiUtilCore.compareElementsByPosition(u1.getElement(), u2.getElement()));
  59   }
  60
  61   public final void moveElement() {
  62     final PsiElement oldElementBody = PyMoveModuleMembersHelper.expandNamedElementBody(myMovedElement);
  63     final PsiFile sourceFile = myMovedElement.getContainingFile();
  64     if (oldElementBody != null) {
  65       PyClassRefactoringUtil.rememberNamedReferences(oldElementBody);
  66       final PsiElement newElementBody = addElementToFile(oldElementBody);
  67       final PsiNamedElement newElement = PyMoveModuleMembersHelper.extractNamedElement(newElementBody);
  68       assert newElement != null;
  69       for (UsageInfo usage : myUsages) {
  70         final PsiElement usageElement = usage.getElement();
  71         if (usageElement != null) {
  72           updateSingleUsage(usageElement, newElement);
  73         }
  74       }
  75       PyClassRefactoringUtil.restoreNamedReferences(newElementBody, myMovedElement, myAllMovedElements);
  76       deleteElement();
  77       optimizeImports(sourceFile);
  78     }
  79   }
  80
  81   private void deleteElement() {
  82     final PsiElement elementBody = PyMoveModuleMembersHelper.expandNamedElementBody(myMovedElement);
  83     assert elementBody != null;
  84     elementBody.delete();
  85   }
  86
  87   private void optimizeImports(@Nullable PsiFile originalFile) {
  88     for (PsiFile usageFile : myOptimizeImportTargets) {
  89       PyClassRefactoringUtil.optimizeImports(usageFile);
  90     }
  91     if (originalFile != null) {
  92       PyClassRefactoringUtil.optimizeImports(originalFile);
  93     }
  94   }
  95
  96   @NotNull
  97   private PsiElement addElementToFile(@NotNull PsiElement element) {
  98     final PsiElement firstUsage = findFirstTopLevelWithUsageAtDestination();
  99     if (firstUsage != null) {
 100       return myDestinationFile.addBefore(element, firstUsage);
 101     }
 102     else {
 103       return myDestinationFile.add(element);
 104     }
 105   }
 106
 107   @Nullable
 108   private PsiElement findFirstTopLevelWithUsageAtDestination() {
 109     final List<PsiElement> topLevelAtDestination = ContainerUtil.mapNotNull(myUsages, usage -> {
 110       final PsiElement element = usage.getElement();
 111       if (element != null && ScopeUtil.getScopeOwner(element) == myDestinationFile && getImportStatementByElement(element) == null) {
 112         return findTopLevelParent(element);
 113       }
 114       return null;
 115     });
 116     if (topLevelAtDestination.isEmpty()) {
 117       return null;
 118     }
 119     return Collections.min(topLevelAtDestination, (e1, e2) -> PsiUtilCore.compareElementsByPosition(e1, e2));
 120   }
 121
 122   @Nullable
 123   private PsiElement findTopLevelParent(@NotNull PsiElement element) {
 124     return PsiTreeUtil.findFirstParent(element, element1 -> element1.getParent() == myDestinationFile);
 125   }
 126
 127   private void updateSingleUsage(@NotNull PsiElement usage, @NotNull PsiNamedElement newElement) {
 128     final PsiFile usageFile = usage.getContainingFile();
 129     if (belongsToSomeMovedElement(usage)) {
 130       return;
 131     }
 132     if (usage instanceof PyQualifiedExpression) {
 133       final PyQualifiedExpression qualifiedExpr = (PyQualifiedExpression)usage;
 134       if (myMovedElement instanceof PyClass && PyNames.INIT.equals(qualifiedExpr.getName())) {
 135         return;
 136       }
 137       else if (qualifiedExpr.isQualified()) {
 138         insertQualifiedImportAndReplaceReference(newElement, qualifiedExpr);
 139       }
 140       else if (usageFile == myMovedElement.getContainingFile()) {
 141         if (usage.getParent() instanceof PyGlobalStatement) {
 142           myScopeOwnersWithGlobal.add(ScopeUtil.getScopeOwner(usage));
 143           if (((PyGlobalStatement)usage.getParent()).getGlobals().length == 1) {
 144             usage.getParent().delete();
 145           }
 146           else {
 147             usage.delete();
 148           }
 149         }
 150         else if (myScopeOwnersWithGlobal.contains(ScopeUtil.getScopeOwner(usage))) {
 151           insertQualifiedImportAndReplaceReference(newElement, qualifiedExpr);
 152         }
 153         else {
 154           insertImportFromAndReplaceReference(newElement, qualifiedExpr);
 155         }
 156       }
 157       else {
 158         final PyImportStatementBase importStmt = getImportStatementByElement(usage);
 159         if (importStmt != null) {
 160           PyClassRefactoringUtil.updateUnqualifiedImportOfElement(importStmt, newElement);
 161         }
 162         else if (resolvesToLocalStarImport(usage)) {
 163           PyClassRefactoringUtil.insertImport(usage, newElement);
 164           myOptimizeImportTargets.add(usageFile);
 165         }
 166       }
 167     }
 168     else if (usage instanceof PyStringLiteralExpression) {
 169       for (PsiReference ref : usage.getReferences()) {
 170         if (ref instanceof PyDunderAllReference) {
 171           usage.delete();
 172         }
 173         else {
 174           if (ref.isReferenceTo(myMovedElement)) {
 175             ref.bindToElement(newElement);
 176           }
 177         }
 178       }
 179     }
 180   }
 181
 182   private boolean belongsToSomeMovedElement(@NotNull final PsiElement element) {
 183     return ContainerUtil.exists(myAllMovedElements, movedElement -> {
 184       final PsiElement movedElementBody = PyMoveModuleMembersHelper.expandNamedElementBody((PsiNamedElement)movedElement);
 185       return PsiTreeUtil.isAncestor(movedElementBody, element, false);
 186     });
 187   }
 188
 189
 190   /**
 191    * <pre><code>
 192    *   print(foo.bar)
 193    * </code></pre>
 194    * is transformed to
 195    * <pre><code>
 196    *   from new import bar
 197    *   print(bar)
 198    * </code></pre>
 199    */
 200   private static void insertImportFromAndReplaceReference(@NotNull PsiNamedElement targetElement,
 201                                                           @NotNull PyQualifiedExpression expression) {
 202     PyClassRefactoringUtil.insertImport(expression, targetElement, null, true);
 203     final PyElementGenerator generator = PyElementGenerator.getInstance(expression.getProject());
 204     final PyExpression generated = generator.createExpressionFromText(LanguageLevel.forElement(expression), expression.getReferencedName());
 205     expression.replace(generated);
 206   }
 207
 208   /**
 209    * <pre><code>
 210    *   print(foo.bar)
 211    * </code></pre>
 212    * is transformed to
 213    * <pre><code>
 214    *   import new
 215    *   print(new.bar)
 216    * </code></pre>
 217    */
 218   private static void insertQualifiedImportAndReplaceReference(@NotNull PsiNamedElement targetElement,
 219                                                                @NotNull PyQualifiedExpression expression) {
 220     final PsiFile file = targetElement.getContainingFile();
 221     final QualifiedName qualifier = QualifiedNameFinder.findCanonicalImportPath(file, expression);
 222     PyClassRefactoringUtil.insertImport(expression, file, null, false);
 223     final PyElementGenerator generator = PyElementGenerator.getInstance(expression.getProject());
 224     final PyExpression generated = generator.createExpressionFromText(LanguageLevel.forElement(expression),
 225                                                                       qualifier + "." + expression.getReferencedName());
 226     expression.replace(generated);
 227   }
 228
 229   private static boolean resolvesToLocalStarImport(@NotNull PsiElement usage) {
 230     // Don't use PyUtil#multiResolveTopPriority here since it filters out low priority ImportedResolveResults
 231     final List<PsiElement> resolvedElements = new ArrayList<>();
 232     if (usage instanceof PyReferenceOwner) {
 233       final PsiPolyVariantReference reference = ((PyReferenceOwner)usage).getReference(PyResolveContext.defaultContext());
 234       for (ResolveResult result : reference.multiResolve(false)) {
 235         resolvedElements.add(result.getElement());
 236       }
 237     }
 238     else {
 239       final PsiReference ref = usage.getReference();
 240       if (ref != null) {
 241         resolvedElements.add(ref.resolve());
 242       }
 243     }
 244     final PsiFile containingFile = usage.getContainingFile();
 245     if (containingFile != null) {
 246       for (PsiElement resolved : resolvedElements) {
 247         if (resolved instanceof PyStarImportElement && resolved.getContainingFile() == containingFile) {
 248           return true;
 249         }
 250       }
 251     }
 252     return false;
 253   }
 254 }