PY-18970 Check that imports in "project" group don't belong to any library root
[idea/community.git] / python / src / com / jetbrains / python / codeInsight / imports / AddImportHelper.java
1 /*
2  * Copyright 2000-2014 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.codeInsight.imports;
17
18 import com.intellij.lang.injection.InjectedLanguageManager;
19 import com.intellij.openapi.diagnostic.Logger;
20 import com.intellij.openapi.module.Module;
21 import com.intellij.openapi.module.ModuleUtilCore;
22 import com.intellij.openapi.projectRoots.Sdk;
23 import com.intellij.openapi.roots.ProjectFileIndex;
24 import com.intellij.openapi.roots.ProjectRootManager;
25 import com.intellij.openapi.util.text.StringUtil;
26 import com.intellij.openapi.vfs.VirtualFile;
27 import com.intellij.psi.*;
28 import com.intellij.psi.util.PsiTreeUtil;
29 import com.intellij.psi.util.QualifiedName;
30 import com.intellij.util.ArrayUtil;
31 import com.intellij.util.IncorrectOperationException;
32 import com.intellij.util.containers.ContainerUtil;
33 import com.jetbrains.python.codeInsight.PyCodeInsightSettings;
34 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
35 import com.jetbrains.python.formatter.PyBlock;
36 import com.jetbrains.python.psi.*;
37 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
38 import com.jetbrains.python.sdk.PythonSdkType;
39 import org.jetbrains.annotations.NotNull;
40 import org.jetbrains.annotations.Nullable;
41
42 import java.util.ArrayList;
43 import java.util.Comparator;
44 import java.util.List;
45
46 import static com.jetbrains.python.psi.PyUtil.as;
47 import static com.jetbrains.python.psi.PyUtil.sure;
48
49 /**
50  * Does the actual job of adding an import statement into a file.
51  * User: dcheryasov
52  * Date: Apr 24, 2009 3:17:59 AM
53  */
54 public class AddImportHelper {
55   private static final Logger LOG = Logger.getInstance("#" + AddImportHelper.class.getName());
56
57   public static final Comparator<PyImportStatementBase> IMPORT_TYPE_THEN_NAME_COMPARATOR = new Comparator<PyImportStatementBase>() {
58     @Override
59     public int compare(@NotNull PyImportStatementBase import1, @NotNull PyImportStatementBase import2) {
60       // normal imports go first, then "from" imports
61       if (import1 instanceof PyImportStatement && import2 instanceof PyFromImportStatement) {
62         return -1;
63       }
64       if (import1 instanceof PyFromImportStatement && import2 instanceof PyImportStatement) {
65         return 1;
66       }
67
68       return ContainerUtil.compareLexicographically(getSortNames(import1), getSortNames(import2));
69     }
70
71     @NotNull
72     public List<String> getSortNames(@NotNull PyImportStatementBase importStatement) {
73       final List<String> result = new ArrayList<String>();
74       final PyFromImportStatement fromImport = as(importStatement, PyFromImportStatement.class);
75       if (fromImport != null) {
76         final QualifiedName source = fromImport.getImportSourceQName();
77         // because of that relative imports go to the end of an import block
78         result.add(StringUtil.repeatSymbol('.', fromImport.getRelativeLevel()));
79         result.add(source != null ? source.toString() : "");
80         if (fromImport.isStarImport()) {
81           result.add("*");
82         }
83       }
84       
85       for (PyImportElement importElement : importStatement.getImportElements()) {
86         final QualifiedName qualifiedName = importElement.getImportedQName();
87         result.add(qualifiedName != null ? qualifiedName.toString() : "");
88       }
89       return result;
90     }
91   };
92
93   public enum ImportPriority {
94     FUTURE,
95     BUILTIN,
96     THIRD_PARTY,
97     PROJECT
98   }
99   
100   private static final ImportPriority UNRESOLVED_SYMBOL_PRIORITY = ImportPriority.THIRD_PARTY;
101
102   private AddImportHelper() {
103   }
104
105   public static void addLocalImportStatement(@NotNull PsiElement element, @NotNull String name) {
106     final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject());
107     final LanguageLevel languageLevel = LanguageLevel.forElement(element);
108
109     final PsiElement anchor = getLocalInsertPosition(element);
110     final PsiElement parentElement = sure(anchor).getParent();
111     if (parentElement != null) {
112       parentElement.addBefore(generator.createImportStatement(languageLevel, name, null), anchor);
113     }
114   }
115
116   public static void addLocalFromImportStatement(@NotNull PsiElement element, @NotNull String qualifier, @NotNull String name) {
117     final PyElementGenerator generator = PyElementGenerator.getInstance(element.getProject());
118     final LanguageLevel languageLevel = LanguageLevel.forElement(element);
119
120     final PsiElement anchor = getLocalInsertPosition(element);
121     final PsiElement parentElement = sure(anchor).getParent();
122     if (parentElement != null) {
123       parentElement.addBefore(generator.createFromImportStatement(languageLevel, qualifier, name, null), anchor);
124     }
125
126   }
127
128   @Nullable
129   public static PsiElement getLocalInsertPosition(@NotNull PsiElement anchor) {
130     return PsiTreeUtil.getParentOfType(anchor, PyStatement.class, false);
131   }
132
133   @Nullable
134   public static PsiElement getFileInsertPosition(final PsiFile file) {
135     return getInsertPosition(file, null, null);
136   }
137
138   @Nullable
139   private static PsiElement getInsertPosition(@NotNull PsiElement insertParent,
140                                               @Nullable PyImportStatementBase newImport,
141                                               @Nullable ImportPriority priority) {
142     PsiElement feeler = insertParent.getFirstChild();
143     if (feeler == null) return null;
144     // skip initial comments and whitespace and try to get just below the last import stmt
145     boolean skippedOverImports = false;
146     boolean skippedOverDoc = false;
147     PsiElement seeker = feeler;
148     final boolean isInjected = InjectedLanguageManager.getInstance(feeler.getProject()).isInjectedFragment(feeler.getContainingFile());
149     PyImportStatementBase importAbove = null, importBelow = null;
150     do {
151       if (feeler instanceof PyImportStatementBase && !isInjected) {
152         final PyImportStatementBase existingImport = (PyImportStatementBase)feeler;
153         if (priority != null && newImport != null) {
154           if (shouldInsertBefore(newImport, existingImport, priority)) {
155             importBelow = existingImport;
156             break;
157           }
158           else {
159             importAbove = existingImport;
160           }
161         }
162         seeker = feeler;
163         feeler = feeler.getNextSibling();
164         skippedOverImports = true;
165       }
166       else if (PyUtil.instanceOf(feeler, PsiWhiteSpace.class, PsiComment.class)) {
167         seeker = feeler;
168         feeler = feeler.getNextSibling();
169       }
170       // maybe we arrived at the doc comment stmt; skip over it, too
171       else if (!skippedOverImports && !skippedOverDoc && insertParent instanceof PyFile) {
172         // this gives the literal; its parent is the expr seeker may have encountered
173         final PsiElement docElem = DocStringUtil.findDocStringExpression((PyElement)insertParent);
174         if (docElem != null && docElem.getParent() == feeler) {
175           feeler = feeler.getNextSibling();
176           seeker = feeler; // skip over doc even if there's nothing below it
177           skippedOverDoc = true;
178         }
179         else {
180           break; // not a doc comment, stop on it
181         }
182       }
183       else {
184         break; // some other statement, stop
185       }
186     }
187     while (feeler != null);
188     final ImportPriority priorityAbove = importAbove != null ? getImportPriority(importAbove) : null;
189     final ImportPriority priorityBelow = importBelow != null ? getImportPriority(importBelow) : null;
190     if (newImport != null && (priorityAbove == null || priorityAbove.compareTo(priority) < 0)) {
191       newImport.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
192     }
193     if (priorityBelow != null) {
194       // actually not necessary because existing import with higher priority (i.e. lower import group) 
195       // probably should have IMPORT_GROUP_BEGIN flag already, but we add it anyway just for safety
196       if (priorityBelow.compareTo(priority) > 0) {
197         importBelow.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, true);
198       }
199       else if (priorityBelow == priority) {
200         importBelow.putCopyableUserData(PyBlock.IMPORT_GROUP_BEGIN, null);
201       }
202     }
203     return seeker;
204   }
205
206   private static boolean shouldInsertBefore(@Nullable PyImportStatementBase newImport,
207                                             @NotNull PyImportStatementBase existingImport,
208                                             @NotNull ImportPriority priority) {
209     final ImportPriority existingImportPriority = getImportPriority(existingImport);
210     final int byPriority = priority.compareTo(existingImportPriority);
211     if (byPriority != 0) {
212       return byPriority < 0;
213     }
214     if (newImport == null) {
215       return false;
216     }
217     return IMPORT_TYPE_THEN_NAME_COMPARATOR.compare(newImport, existingImport) < 0;
218   }
219
220   @NotNull
221   public static ImportPriority getImportPriority(@NotNull PyImportStatementBase importStatement) {
222     final PsiElement resolved;
223     if (importStatement instanceof PyFromImportStatement) {
224       final PyFromImportStatement fromImportStatement = (PyFromImportStatement)importStatement;
225       if (fromImportStatement.isFromFuture()) {
226         return ImportPriority.FUTURE;
227       }
228       resolved = fromImportStatement.resolveImportSource();
229     }
230     else {
231       final PyImportElement firstImportElement = ArrayUtil.getFirstElement(importStatement.getImportElements());
232       if (firstImportElement == null) {
233         return UNRESOLVED_SYMBOL_PRIORITY;
234       }
235       resolved = firstImportElement.resolve();
236     }
237     if (resolved == null) {
238       return UNRESOLVED_SYMBOL_PRIORITY;
239     }
240
241     final PsiFileSystemItem resolvedFileOrDir;
242     if (resolved instanceof PsiDirectory) {
243       resolvedFileOrDir = (PsiFileSystemItem)resolved;
244     }
245     // resolved symbol may be PsiPackage in Jython
246     else if (resolved instanceof PsiDirectoryContainer) {
247       resolvedFileOrDir = ArrayUtil.getFirstElement(((PsiDirectoryContainer)resolved).getDirectories());
248     }
249     else {
250       resolvedFileOrDir = resolved.getContainingFile();
251     }
252     
253     if (resolvedFileOrDir == null) {
254       return UNRESOLVED_SYMBOL_PRIORITY;
255     }
256     
257     return getImportPriority(importStatement, resolvedFileOrDir);
258   }
259
260   @NotNull
261   public static ImportPriority getImportPriority(@NotNull PsiElement importLocation, @NotNull PsiFileSystemItem toImport) {
262     final VirtualFile vFile = toImport.getVirtualFile();
263     if (vFile == null) {
264       return UNRESOLVED_SYMBOL_PRIORITY;
265     }
266     final ProjectRootManager projectRootManager = ProjectRootManager.getInstance(toImport.getProject());
267     final ProjectFileIndex fileIndex = projectRootManager.getFileIndex();
268     if (fileIndex.isInContent(vFile) && !fileIndex.isInLibraryClasses(vFile)) {
269       return ImportPriority.PROJECT;
270     }
271     final Module module = ModuleUtilCore.findModuleForPsiElement(importLocation);
272     final Sdk pythonSdk = module != null ? PythonSdkType.findPythonSdk(module) : projectRootManager.getProjectSdk();
273
274     return PythonSdkType.isStdLib(vFile, pythonSdk) ? ImportPriority.BUILTIN : ImportPriority.THIRD_PARTY;
275   }
276
277   /**
278    * Adds an import statement, if it doesn't exist yet, presumably below all other initial imports in the file.
279    *
280    * @param file   where to operate
281    * @param name   which to import (qualified is OK)
282    * @param asName optional name for 'as' clause
283    * @param anchor place where the imported name was used. It will be used to determine proper block where new import should be inserted,
284    *               e.g. inside conditional block or try/except statement. Also if anchor is another import statement, new import statement
285    *               will be inserted right after it.
286    * @return whether import statement was actually added
287    */
288   public static boolean addImportStatement(@NotNull PsiFile file,
289                                            @NotNull String name,
290                                            @Nullable String asName,
291                                            @Nullable ImportPriority priority,
292                                            @Nullable PsiElement anchor) {
293     if (!(file instanceof PyFile)) {
294       return false;
295     }
296     final List<PyImportElement> existingImports = ((PyFile)file).getImportTargets();
297     for (PyImportElement element : existingImports) {
298       final QualifiedName qName = element.getImportedQName();
299       if (qName != null && name.equals(qName.toString())) {
300         if ((asName != null && asName.equals(element.getAsName())) || asName == null) {
301           return false;
302         }
303       }
304     }
305
306     final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
307     final LanguageLevel languageLevel = LanguageLevel.forElement(file);
308     final PyImportStatement importNodeToInsert = generator.createImportStatement(languageLevel, name, asName);
309     final PyImportStatementBase importStatement = PsiTreeUtil.getParentOfType(anchor, PyImportStatementBase.class, false);
310     final PsiElement insertParent = importStatement != null && importStatement.getContainingFile() == file ?
311                                     importStatement.getParent() : file;
312     try {
313       if (anchor instanceof PyImportStatementBase) {
314         insertParent.addAfter(importNodeToInsert, anchor);
315       }
316       else {
317         insertParent.addBefore(importNodeToInsert, getInsertPosition(insertParent, importNodeToInsert, priority));
318       }
319     }
320     catch (IncorrectOperationException e) {
321       LOG.error(e);
322     }
323     return true;
324   }
325
326   /**
327    * Adds a new {@link PyFromImportStatement} statement within other top-level imports or as specified by anchor.
328    *
329    * @param file   where to operate
330    * @param from   import source (reference after {@code from} keyword)
331    * @param name   imported name (identifier after {@code import} keyword)
332    * @param asName optional alias (identifier after {@code as} keyword)
333    * @param anchor place where the imported name was used. It will be used to determine proper block where new import should be inserted,
334    *               e.g. inside conditional block or try/except statement. Also if anchor is another import statement, new import statement
335    *               will be inserted right after it.
336    * @see #addOrUpdateFromImportStatement
337    */
338   public static void addFromImportStatement(@NotNull PsiFile file,
339                                             @NotNull String from,
340                                             @NotNull String name,
341                                             @Nullable String asName,
342                                             @Nullable ImportPriority priority,
343                                             @Nullable PsiElement anchor) {
344     final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
345     final LanguageLevel languageLevel = LanguageLevel.forElement(file);
346     final PyFromImportStatement newImport = generator.createFromImportStatement(languageLevel, from, name, asName);
347     addFromImportStatement(file, newImport, priority, anchor);
348   }
349
350   /**
351    * Adds a new {@link PyFromImportStatement} statement within other top-level imports or as specified by anchor.
352    *
353    * @param file      where to operate
354    * @param newImport new "from import" statement to insert. It may be generated, because it won't be used for resolving anyway. 
355    *                  You might want to use overloaded version of this method to generate such statement automatically.
356    * @param anchor    place where the imported name was used. It will be used to determine proper block where new import should be inserted,
357    *                  e.g. inside conditional block or try/except statement. Also if anchor is another import statement, new import statement
358    *                  will be inserted right after it.
359    * @see #addFromImportStatement(PsiFile, String, String, String, ImportPriority, PsiElement)
360    * @see #addFromImportStatement
361    */
362   public static void addFromImportStatement(@NotNull PsiFile file,
363                                             @NotNull PyFromImportStatement newImport,
364                                             @Nullable ImportPriority priority,
365                                             @Nullable PsiElement anchor) {
366     try {
367       final PyImportStatementBase parentImport = PsiTreeUtil.getParentOfType(anchor, PyImportStatementBase.class, false);
368       final PsiElement insertParent;
369       if (parentImport != null && parentImport.getContainingFile() == file) {
370         insertParent = parentImport.getParent();
371       }
372       else {
373         insertParent = file;
374       }
375       if (InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) {
376         final PsiElement element = insertParent.addBefore(newImport, getInsertPosition(insertParent, newImport, priority));
377         PsiElement whitespace = element.getNextSibling();
378         if (!(whitespace instanceof PsiWhiteSpace)) {
379           whitespace = PsiParserFacade.SERVICE.getInstance(file.getProject()).createWhiteSpaceFromText("  >>> ");
380         }
381         insertParent.addBefore(whitespace, element);
382       }
383       else {
384         if (anchor instanceof PyImportStatementBase) {
385           insertParent.addAfter(newImport, anchor);
386         }
387         else {
388           insertParent.addBefore(newImport, getInsertPosition(insertParent, newImport, priority));
389         }
390       }
391     }
392     catch (IncorrectOperationException e) {
393       LOG.error(e);
394     }
395   }
396
397   /**
398    * Adds new {@link PyFromImportStatement} in file or append {@link PyImportElement} to
399    * existing from import statement.
400    *
401    * @param file     module where import will be added
402    * @param from     import source (reference after {@code from} keyword)
403    * @param name     imported name (identifier after {@code import} keyword)
404    * @param asName   optional alias (identifier after {@code as} keyword)
405    * @param priority optional import priority used to sort imports
406    * @param anchor   place where the imported name was used. It will be used to determine proper block where new import should be inserted,
407    *                 e.g. inside conditional block or try/except statement. Also if anchor is another import statement, new import statement
408    *                 will be inserted right after it.
409    * @return whether import was actually added
410    * @see #addFromImportStatement
411    */
412   public static boolean addOrUpdateFromImportStatement(@NotNull PsiFile file,
413                                                        @NotNull String from,
414                                                        @NotNull String name,
415                                                        @Nullable String asName,
416                                                        @Nullable ImportPriority priority,
417                                                        @Nullable PsiElement anchor) {
418     final List<PyFromImportStatement> existingImports = ((PyFile)file).getFromImports();
419     for (PyFromImportStatement existingImport : existingImports) {
420       if (existingImport.isStarImport()) {
421         continue;
422       }
423       final QualifiedName qName = existingImport.getImportSourceQName();
424       if (qName != null && qName.toString().equals(from) && existingImport.getRelativeLevel() == 0) {
425         for (PyImportElement el : existingImport.getImportElements()) {
426           final QualifiedName importedQName = el.getImportedQName();
427           if (importedQName != null && StringUtil.equals(name, importedQName.toString()) && StringUtil.equals(asName, el.getAsName())) {
428             return false;
429           }
430         }
431         final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
432         final PyImportElement importElement = generator.createImportElement(LanguageLevel.forElement(file), name);
433         existingImport.add(importElement);
434         return false;
435       }
436     }
437     addFromImportStatement(file, from, name, asName, priority, anchor);
438     return true;
439   }
440
441   /**
442    * Adds either {@link PyFromImportStatement} or {@link PyImportStatement}
443    * to specified target depending on user preferences and whether it's possible to import element via "from" form of import
444    * (e.g. consider top level module).
445    *
446    * @param target  element import is pointing to
447    * @param file    file where import will be inserted
448    * @param element used to determine where to insert import
449    * @see PyCodeInsightSettings#PREFER_FROM_IMPORT
450    * @see #addImportStatement
451    * @see #addOrUpdateFromImportStatement
452    */
453   public static void addImport(final PsiNamedElement target, final PsiFile file, final PyElement element) {
454     final boolean useQualified = !PyCodeInsightSettings.getInstance().PREFER_FROM_IMPORT;
455     final PsiFileSystemItem toImport =
456       target instanceof PsiFileSystemItem ? ((PsiFileSystemItem)target).getParent() : target.getContainingFile();
457     if (toImport == null) return;
458     final ImportPriority priority = getImportPriority(file, toImport);
459     final QualifiedName qName = QualifiedNameFinder.findCanonicalImportPath(target, element);
460     if (qName == null) return;
461     String path = qName.toString();
462     if (target instanceof PsiFileSystemItem && qName.getComponentCount() == 1) {
463       addImportStatement(file, path, null, priority, element);
464     }
465     else {
466       final QualifiedName toImportQName = QualifiedNameFinder.findCanonicalImportPath(toImport, element);
467       if (toImportQName == null) return;
468       if (useQualified) {
469         addImportStatement(file, path, null, priority, element);
470         final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(file.getProject());
471         final String targetName = PyUtil.getElementNameWithoutExtension(target);
472         element.replace(elementGenerator.createExpressionFromText(LanguageLevel.forElement(target), toImportQName + "." + targetName));
473       }
474       else {
475         final String name = target.getName();
476         if (name != null)
477           addOrUpdateFromImportStatement(file, toImportQName.toString(), name, null, priority, element);
478       }
479     }
480   }
481 }