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 }