python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.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.psi.resolve;
  17
  18 import com.intellij.openapi.extensions.Extensions;
  19 import com.intellij.openapi.fileTypes.ExtensionFileNameMatcher;
  20 import com.intellij.openapi.fileTypes.FileNameMatcher;
  21 import com.intellij.openapi.fileTypes.FileTypeManager;
  22 import com.intellij.openapi.module.Module;
  23 import com.intellij.openapi.module.ModuleUtilCore;
  24 import com.intellij.openapi.projectRoots.Sdk;
  25 import com.intellij.openapi.roots.FileIndexFacade;
  26 import com.intellij.openapi.util.io.FileUtil;
  27 import com.intellij.openapi.vfs.VirtualFile;
  28 import com.intellij.psi.PsiDirectory;
  29 import com.intellij.psi.PsiElement;
  30 import com.intellij.psi.PsiFile;
  31 import com.intellij.psi.PsiInvalidElementAccessException;
  32 import com.intellij.psi.util.PsiTreeUtil;
  33 import com.intellij.psi.util.QualifiedName;
  34 import com.intellij.util.containers.HashSet;
  35 import com.jetbrains.python.PyNames;
  36 import com.jetbrains.python.PythonFileType;
  37 import com.jetbrains.python.psi.*;
  38 import com.jetbrains.python.psi.impl.*;
  39 import com.jetbrains.python.psi.types.PyModuleType;
  40 import com.jetbrains.python.psi.types.PyType;
  41 import org.jetbrains.annotations.NotNull;
  42 import org.jetbrains.annotations.Nullable;
  43
  44 import java.util.ArrayList;
  45 import java.util.Collections;
  46 import java.util.List;
  47 import java.util.Set;
  48
  49 import static com.jetbrains.python.psi.FutureFeature.ABSOLUTE_IMPORT;
  50
  51 /**
  52  * @author dcheryasov
  53  */
  54 public class ResolveImportUtil {
  55   private ResolveImportUtil() {
  56   }
  57
  58   private static final ThreadLocal<Set<String>> ourBeingImported = new ThreadLocal<Set<String>>() {
  59     @Override
  60     protected Set<String> initialValue() {
  61       return new HashSet<String>();
  62     }
  63   };
  64
  65   public static boolean isAbsoluteImportEnabledFor(PsiElement foothold) {
  66     if (foothold != null) {
  67       PsiFile file = foothold.getContainingFile();
  68       if (file instanceof PyFile) {
  69         final PyFile pyFile = (PyFile)file;
  70         if (pyFile.getLanguageLevel().isPy3K()) {
  71           return true;
  72         }
  73         return pyFile.hasImportFromFuture(ABSOLUTE_IMPORT);
  74       }
  75     }
  76     // if the relevant import is below the foothold, it is either legal or we've detected the offending statement already
  77     return false;
  78   }
  79
  80
  81   /**
  82    * Finds a directory that many levels above a given file, making sure that every level has an __init__.py.
  83    *
  84    * @param base  file that works as a reference.
  85    * @param depth must be positive, 1 means the dir that contains base, 2 is one dir above, etc.
  86    * @return found directory, or null.
  87    */
  88   @Nullable
  89   public static PsiDirectory stepBackFrom(PsiFile base, int depth) {
  90     if (depth == 0) {
  91       return base.getContainingDirectory();
  92     }
  93     PsiDirectory result;
  94     if (base != null) {
  95       base = base.getOriginalFile(); // just to make sure
  96       result = base.getContainingDirectory();
  97       int count = 1;
  98       while (result != null && PyUtil.isPackage(result, base)) {
  99         if (count >= depth) return result;
 100         result = result.getParentDirectory();
 101         count += 1;
 102       }
 103     }
 104     return null;
 105   }
 106
 107   @Nullable
 108   public static PsiElement resolveImportElement(PyImportElement importElement, @NotNull final QualifiedName qName) {
 109     List<RatedResolveResult> targets;
 110     final PyStatement importStatement = importElement.getContainingImportStatement();
 111     if (importStatement instanceof PyFromImportStatement) {
 112       targets = resolveNameInFromImport((PyFromImportStatement)importStatement, qName);
 113     }
 114     else { // "import foo"
 115       targets = resolveNameInImportStatement(importElement, qName);
 116     }
 117     final List<RatedResolveResult> resultList = RatedResolveResult.sorted(targets);
 118     return resultList.size() > 0 ? resultList.get(0).getElement() : null;
 119   }
 120
 121   public static List<RatedResolveResult> resolveNameInImportStatement(PyImportElement importElement, @NotNull QualifiedName qName) {
 122     final PsiFile file = importElement.getContainingFile().getOriginalFile();
 123     boolean absoluteImportEnabled = isAbsoluteImportEnabledFor(importElement);
 124     final List<PsiElement> modules = resolveModule(qName, file, absoluteImportEnabled, 0);
 125     return rateResults(modules);
 126   }
 127
 128   public static List<RatedResolveResult> resolveNameInFromImport(PyFromImportStatement importStatement, @NotNull QualifiedName qName) {
 129     PsiFile file = importStatement.getContainingFile().getOriginalFile();
 130     String name = qName.getComponents().get(0);
 131
 132     final List<PsiElement> candidates = importStatement.resolveImportSourceCandidates();
 133     List<PsiElement> resultList = new ArrayList<PsiElement>();
 134     for (PsiElement candidate : candidates) {
 135       if (!candidate.isValid()) {
 136         throw new PsiInvalidElementAccessException(candidate, "Got an invalid candidate from resolveImportSourceCandidates(): " + candidate.getClass());
 137       }
 138       if (candidate instanceof PsiDirectory) {
 139         candidate = PyUtil.getPackageElement((PsiDirectory)candidate, importStatement);
 140       }
 141       PsiElement result = resolveChild(candidate, name, file, false, true);
 142       if (result != null) {
 143         if (!result.isValid()) {
 144           throw new PsiInvalidElementAccessException(result, "Got an invalid candidate from resolveChild(): " + result.getClass());
 145         }
 146         resultList.add(result);
 147       }
 148     }
 149     if (!resultList.isEmpty()) {
 150       return rateResults(resultList);
 151     }
 152     return Collections.emptyList();
 153   }
 154
 155   @NotNull
 156   public static List<PsiElement> resolveFromImportStatementSource(@NotNull PyFromImportStatement fromImportStatement,
 157                                                                   @Nullable QualifiedName qName) {
 158     final boolean absoluteImportEnabled = isAbsoluteImportEnabledFor(fromImportStatement);
 159     final PsiFile file = fromImportStatement.getContainingFile();
 160     return resolveModule(qName, file, absoluteImportEnabled, fromImportStatement.getRelativeLevel());
 161   }
 162
 163   /**
 164    * Resolves a module reference in a general case.
 165    *
 166    *
 167    * @param qualifiedName     qualified name of the module reference to resolve
 168    * @param sourceFile        where that reference resides; serves as PSI foothold to determine module, project, etc.
 169    * @param importIsAbsolute  if false, try old python 2.x's "relative first, absolute next" approach.
 170    * @param relativeLevel     if > 0, step back from sourceFile and resolve from there (even if importIsAbsolute is false!).
 171    * @return list of possible candidates
 172    */
 173   @NotNull
 174   public static List<PsiElement> resolveModule(@Nullable QualifiedName qualifiedName, @Nullable PsiFile sourceFile,
 175                                                boolean importIsAbsolute, int relativeLevel) {
 176     if (qualifiedName == null || sourceFile == null) {
 177       return Collections.emptyList();
 178     }
 179     final String marker = qualifiedName + "#" + Integer.toString(relativeLevel);
 180     final Set<String> beingImported = ourBeingImported.get();
 181     if (beingImported.contains(marker)) {
 182       return Collections.emptyList(); // break endless loop in import
 183     }
 184     try {
 185       beingImported.add(marker);
 186       final QualifiedNameResolver visitor = new QualifiedNameResolverImpl(qualifiedName).fromElement(sourceFile);
 187       if (relativeLevel > 0) {
 188         // "from ...module import"
 189         visitor.withRelative(relativeLevel).withoutRoots();
 190       }
 191       else {
 192         // "from module import"
 193         if (!importIsAbsolute) {
 194           visitor.withRelative(0);
 195         }
 196       }
 197       List<PsiElement> results = visitor.resultsAsList();
 198       if (results.isEmpty() && relativeLevel == 0 && !importIsAbsolute) {
 199         results = resolveRelativeImportAsAbsolute(sourceFile, qualifiedName);
 200       }
 201       return results;
 202     }
 203     finally {
 204       beingImported.remove(marker);
 205     }
 206   }
 207
 208   /**
 209    * Try to resolve relative import as absolute in roots, not in its parent directory.
 210    *
 211    * This may be useful for resolving to child skeleton modules located in other directories.
 212    *
 213    * @param foothold        foothold file.
 214    * @param qualifiedName   relative import name.
 215    * @return                list of resolved elements.
 216    */
 217   @NotNull
 218   private static List<PsiElement> resolveRelativeImportAsAbsolute(@NotNull PsiFile foothold,
 219                                                                   @NotNull QualifiedName qualifiedName) {
 220     final VirtualFile virtualFile = foothold.getVirtualFile();
 221     if (virtualFile == null) return Collections.emptyList();
 222     final boolean inSource = FileIndexFacade.getInstance(foothold.getProject()).isInContent(virtualFile);
 223     if (inSource) return Collections.emptyList();
 224     final PsiDirectory containingDirectory = foothold.getContainingDirectory();
 225     if (containingDirectory != null) {
 226       final QualifiedName containingPath = QualifiedNameFinder.findCanonicalImportPath(containingDirectory, null);
 227       if (containingPath != null && containingPath.getComponentCount() > 0) {
 228         final QualifiedName absolutePath = containingPath.append(qualifiedName.toString());
 229         final QualifiedNameResolver absoluteVisitor = new QualifiedNameResolverImpl(absolutePath).fromElement(foothold);
 230         return absoluteVisitor.resultsAsList();
 231       }
 232     }
 233     return Collections.emptyList();
 234   }
 235
 236   @Nullable
 237   public static PsiElement resolveModuleInRoots(@NotNull QualifiedName moduleQualifiedName, @Nullable PsiElement foothold) {
 238     if (foothold == null) return null;
 239     QualifiedNameResolver visitor = new QualifiedNameResolverImpl(moduleQualifiedName).fromElement(foothold);
 240     return visitor.firstResult();
 241   }
 242
 243   @Nullable
 244   static PythonPathCache getPathCache(PsiElement foothold) {
 245     PythonPathCache cache = null;
 246     final Module module = ModuleUtilCore.findModuleForPsiElement(foothold);
 247     if (module != null) {
 248       cache = PythonModulePathCache.getInstance(module);
 249     }
 250     else {
 251       final Sdk sdk = PyBuiltinCache.findSdkForFile(foothold.getContainingFile());
 252       if (sdk != null) {
 253         cache = PythonSdkPathCache.getInstance(foothold.getProject(), sdk);
 254       }
 255     }
 256     return cache;
 257   }
 258
 259   /**
 260    * Tries to find referencedName under the parent element.
 261    *
 262    * @param parent          element under which to look for referenced name; if null, null is returned.
 263    * @param referencedName  which name to look for.
 264    * @param containingFile  where we're in.
 265    * @param fileOnly        if true, considers only a PsiFile child as a valid result; non-file hits are ignored.
 266    * @param checkForPackage if true, directories are returned only if they contain __init__.py
 267    * @return the element the referencedName resolves to, or null.
 268    */
 269   @Nullable
 270   public static PsiElement resolveChild(@Nullable final PsiElement parent, @NotNull final String referencedName,
 271                                         @Nullable final PsiFile containingFile, boolean fileOnly, boolean checkForPackage) {
 272     PsiDirectory dir = null;
 273     PsiElement resultElement = null;
 274     final PyResolveContext resolveContext = PyResolveContext.defaultContext();
 275     if (parent instanceof PyFileImpl) {
 276       PsiElement possibleResult = null;
 277       if (PyNames.INIT_DOT_PY.equals(((PyFile)parent).getName())) {
 278         dir = ((PyFile)parent).getContainingDirectory();
 279         possibleResult = resolveInDirectory(referencedName, containingFile, dir, fileOnly, checkForPackage);
 280       }
 281
 282       final PyModuleType moduleType = new PyModuleType((PyFile)parent);
 283       final List<? extends RatedResolveResult> results = moduleType.resolveMember(referencedName, null, AccessDirection.READ,
 284                                                                                   resolveContext);
 285       final PsiElement moduleMember = results != null && !results.isEmpty() ? results.get(0).getElement() : null;
 286       if (!fileOnly || PyUtil.instanceOf(moduleMember, PsiFile.class, PsiDirectory.class)) {
 287         resultElement = moduleMember;
 288       }
 289       if (resultElement != null && !PyUtil.instanceOf(resultElement, PsiFile.class, PsiDirectory.class) &&
 290           PsiTreeUtil.getStubOrPsiParentOfType(resultElement, PyExceptPart.class) == null && !isDunderAll(resultElement)) {
 291         return resultElement;
 292       }
 293
 294       if (possibleResult != null) return possibleResult;
 295
 296       if (resultElement != null) {
 297         return resultElement;
 298       }
 299     }
 300     else if (parent instanceof PsiDirectory) {
 301       dir = (PsiDirectory)parent;
 302     }
 303     else if (parent != null) {
 304       PyType refType = PyReferenceExpressionImpl.getReferenceTypeFromProviders(parent, resolveContext.getTypeEvalContext(), null);
 305       if (refType != null) {
 306         final List<? extends RatedResolveResult> result = refType.resolveMember(referencedName, null, AccessDirection.READ, resolveContext);
 307         if (result != null && !result.isEmpty()) {
 308           return result.get(0).getElement();
 309         }
 310       }
 311     }
 312     if (dir != null) {
 313       final PsiElement result = resolveInDirectory(referencedName, containingFile, dir, fileOnly, checkForPackage);
 314       if (result != null) {
 315         return result;
 316       }
 317       if (parent instanceof PsiFile) {
 318         final PsiElement element = new QualifiedNameResolverImpl(referencedName).fromElement(parent).withoutRoots().firstResult();
 319         if (element != null) {
 320           return element;
 321         }
 322       }
 323     }
 324     return null;
 325   }
 326
 327   private static boolean isDunderAll(@NotNull PsiElement element) {
 328     return (element instanceof PyElement) && PyNames.ALL.equals(((PyElement)element).getName());
 329   }
 330
 331   @Nullable
 332   private static PsiElement resolveInDirectory(final String referencedName, @Nullable final PsiFile containingFile,
 333                                                final PsiDirectory dir, boolean isFileOnly, boolean checkForPackage) {
 334     if (referencedName == null) return null;
 335
 336     final PsiDirectory subdir = dir.findSubdirectory(referencedName);
 337     if (subdir != null && (!checkForPackage || PyUtil.isPackage(subdir, containingFile))) {
 338       return subdir;
 339     }
 340
 341     final PsiFile module = findPyFileInDir(dir, referencedName);
 342     if (module != null) return module;
 343
 344     if (!isFileOnly) {
 345       // not a subdir, not a file; could be a name in parent/__init__.py
 346       final PsiFile initPy = dir.findFile(PyNames.INIT_DOT_PY);
 347       if (initPy == containingFile) return null; // don't dive into the file we're in
 348       if (initPy instanceof PyFile) {
 349         return ((PyFile)initPy).getElementNamed(referencedName);
 350       }
 351     }
 352     return null;
 353   }
 354
 355   @Nullable
 356   private static PsiFile findPyFileInDir(PsiDirectory dir, String referencedName) {
 357     PsiFile file = dir.findFile(referencedName + PyNames.DOT_PY);
 358     if (file == null) {
 359       final List<FileNameMatcher> associations = FileTypeManager.getInstance().getAssociations(PythonFileType.INSTANCE);
 360       for (FileNameMatcher association : associations) {
 361         if (association instanceof ExtensionFileNameMatcher) {
 362           file = dir.findFile(referencedName + "." + ((ExtensionFileNameMatcher)association).getExtension());
 363           if (file != null) break;
 364         }
 365       }
 366     }
 367     if (file != null && FileUtil.getNameWithoutExtension(file.getName()).equals(referencedName)) {
 368       return file;
 369     }
 370     return null;
 371   }
 372
 373   public static ResolveResultList rateResults(List<? extends PsiElement> targets) {
 374     ResolveResultList ret = new ResolveResultList();
 375     for (PsiElement target : targets) {
 376       if (target instanceof PsiDirectory) {
 377         target = PyUtil.getPackageElement((PsiDirectory)target, null);
 378       }
 379       if (target != null) {   // Ignore non-package dirs, worthless
 380         int rate = RatedResolveResult.RATE_HIGH;
 381         if (target instanceof PyFile) {
 382           VirtualFile vFile = ((PyFile)target).getVirtualFile();
 383           if (vFile != null && vFile.getLength() > 0) {
 384             rate += 100;
 385           }
 386           for (PyResolveResultRater rater : Extensions.getExtensions(PyResolveResultRater.EP_NAME)) {
 387             rate += rater.getImportElementRate(target);
 388           }
 389         }
 390         ret.poke(target, rate);
 391       }
 392     }
 393     return ret;
 394   }
 395
 396   /**
 397    * @param element what we test (identifier, reference, import element, etc)
 398    * @return the how the element relates to an enclosing import statement, if any
 399    * @see PointInImport
 400    */
 401   @NotNull
 402   public static PointInImport getPointInImport(@NotNull PsiElement element) {
 403     final PsiElement parent = PsiTreeUtil.getNonStrictParentOfType(element, PyImportElement.class, PyFromImportStatement.class);
 404     if (parent instanceof PyFromImportStatement) {
 405       return PointInImport.AS_MODULE; // from foo ...
 406     }
 407     if (parent instanceof PyImportElement) {
 408       final PsiElement statement = parent.getParent();
 409       if (statement instanceof PyImportStatement) {
 410         return PointInImport.AS_MODULE; // import foo,...
 411       }
 412       else if (statement instanceof PyFromImportStatement) {
 413         return PointInImport.AS_NAME;
 414       }
 415     }
 416     return PointInImport.NONE;
 417   }
 418 }