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