f1b9fa82f704ee2e682e7aad5cd8a7e87b9f2fc2
[idea/community.git] / 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 }