9699a8ff149818b1f310b00ab6c7c28f8461944b
[idea/community.git] / python / src / com / jetbrains / python / psi / resolve / QualifiedNameResolverImpl.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.google.common.base.Preconditions;
19 import com.google.common.collect.Lists;
20 import com.google.common.collect.Sets;
21 import com.intellij.facet.Facet;
22 import com.intellij.facet.FacetManager;
23 import com.intellij.openapi.extensions.Extensions;
24 import com.intellij.openapi.module.Module;
25 import com.intellij.openapi.module.ModuleManager;
26 import com.intellij.openapi.project.Project;
27 import com.intellij.openapi.projectRoots.Sdk;
28 import com.intellij.openapi.roots.FileIndexFacade;
29 import com.intellij.openapi.vfs.VirtualFile;
30 import com.intellij.psi.PsiDirectory;
31 import com.intellij.psi.PsiElement;
32 import com.intellij.psi.PsiFile;
33 import com.intellij.psi.PsiNamedElement;
34 import com.intellij.psi.util.PsiTreeUtil;
35 import com.intellij.psi.util.QualifiedName;
36 import com.jetbrains.python.codeInsight.userSkeletons.PyUserSkeletonsUtil;
37 import com.jetbrains.python.console.PydevConsoleRunner;
38 import com.jetbrains.python.facet.PythonPathContributingFacet;
39 import com.jetbrains.python.psi.PyFile;
40 import com.jetbrains.python.psi.impl.PyBuiltinCache;
41 import com.jetbrains.python.psi.impl.PyImportResolver;
42 import com.jetbrains.python.sdk.PySdkUtil;
43 import com.jetbrains.python.sdk.PythonSdkType;
44 import org.jetbrains.annotations.NotNull;
45 import org.jetbrains.annotations.Nullable;
46
47 import java.util.ArrayList;
48 import java.util.Collections;
49 import java.util.List;
50 import java.util.Set;
51
52 /**
53  * Resolves the specified qualified name in the specified context (module, all modules or a file) to a file or directory.
54  *
55  * @author yole
56  */
57 public class QualifiedNameResolverImpl implements RootVisitor, QualifiedNameResolver {
58   boolean myCheckForPackage = true;
59   private final QualifiedNameResolveContext myContext = new QualifiedNameResolveContext();
60   private final @NotNull QualifiedName myQualifiedName;
61   final Set<PsiElement> mySourceResults = Sets.newLinkedHashSet();
62   final Set<PsiElement> myLibResults = Sets.newLinkedHashSet();
63   final Set<PsiElement> myForeignResults = Sets.newLinkedHashSet();
64   private boolean myVisitAllModules = false;
65   private int myRelativeLevel = -1;
66   private boolean myWithoutRoots;
67   private boolean myWithoutForeign;
68   private boolean myWithMembers;
69
70   public QualifiedNameResolverImpl(@NotNull String qNameString) {
71     myQualifiedName = QualifiedName.fromDottedString(qNameString);
72   }
73
74   public QualifiedNameResolverImpl(@NotNull QualifiedName qName) {
75     myQualifiedName = qName;
76   }
77
78   @Override
79   public QualifiedNameResolver withContext(QualifiedNameResolveContext context) {
80     myContext.copyFrom(context);
81     return this;
82   }
83
84   @Override
85   public QualifiedNameResolver fromElement(@NotNull PsiElement foothold) {
86     myContext.setFromElement(foothold);
87     if (PydevConsoleRunner.isInPydevConsole(foothold)) {
88       withAllModules();
89       Sdk sdk = PydevConsoleRunner.getConsoleSdk(foothold);
90       if (sdk != null) {
91         myContext.setSdk(sdk);
92       }
93     }
94     return this;
95   }
96
97   @Override
98   public QualifiedNameResolver fromModule(@NotNull Module module) {
99     myContext.setFromModule(module);
100     return this;
101   }
102
103   @Override
104   public QualifiedNameResolver fromSdk(@NotNull Project project, @NotNull Sdk sdk) {
105     myContext.setFromSdk(project, sdk);
106     return this;
107   }
108
109   private boolean isAcceptRootAsTopLevelPackage() {
110     Module module = myContext.getModule();
111     if (module != null) {
112       Facet[] facets = FacetManager.getInstance(module).getAllFacets();
113       for (Facet facet : facets) {
114         if (facet instanceof PythonPathContributingFacet && ((PythonPathContributingFacet)facet).acceptRootAsTopLevelPackage()) {
115           return true;
116         }
117       }
118     }
119     return false;
120   }
121
122   @Override
123   public QualifiedNameResolver withAllModules() {
124     myVisitAllModules = true;
125     return this;
126   }
127
128   /**
129    * Specifies that we need to look for the name in the specified SDK (instead of the SDK assigned to the module, if any).
130    *
131    * @param sdk the SDK in which the name should be searched.
132    * @return this
133    */
134   @Override
135   public QualifiedNameResolver withSdk(Sdk sdk) {
136     myContext.setSdk(sdk);
137     return this;
138   }
139
140   /**
141    * Specifies whether we should attempt to resolve imports relative to the current file.
142    *
143    * @param relativeLevel if >= 0, we try to resolve at the specified number of levels above the current file.
144    * @return this
145    */
146   @Override
147   public QualifiedNameResolver withRelative(int relativeLevel) {
148     myRelativeLevel = relativeLevel;
149     return this;
150   }
151
152   /**
153    * Specifies that we should only try to resolve relative to the current file, not in roots.
154    *
155    * @return this
156    */
157   @Override
158   public QualifiedNameResolver withoutRoots() {
159     myWithoutRoots = true;
160     return this;
161   }
162
163   @Override
164   public QualifiedNameResolver withoutForeign() {
165     myWithoutForeign = true;
166     return this;
167   }
168
169   @Override
170   public QualifiedNameResolver withMembers() {
171     myWithMembers = true;
172     return this;
173   }
174
175   /**
176    * Specifies that we're looking for a file in a directory hierarchy, not a module in the Python package hierarchy
177    * (so we don't need to check for existence of __init__.py)
178    *
179    * @return
180    */
181   @Override
182   public QualifiedNameResolver withPlainDirectories() {
183     myCheckForPackage = false;
184     return this;
185   }
186
187   public boolean visitRoot(final VirtualFile root, @Nullable Module module, @Nullable Sdk sdk, boolean isModuleSource) {
188     if (!root.isValid()) {
189       return true;
190     }
191     if (root.equals(PyUserSkeletonsUtil.getUserSkeletonsDirectory())) {
192       return true;
193     }
194     PsiElement resolveResult = resolveInRoot(root);
195     if (resolveResult != null) {
196       addRoot(resolveResult, isModuleSource);
197     }
198
199     if (isAcceptRootAsTopLevelPackage() && myQualifiedName.matchesPrefix(QualifiedName.fromDottedString(root.getName()))) {
200       resolveResult = resolveInRoot(root.getParent());
201       if (resolveResult != null) {
202         addRoot(resolveResult, isModuleSource);
203       }
204     }
205
206     return true;
207   }
208
209   private void addRoot(PsiElement resolveResult, boolean isModuleSource) {
210     if (isModuleSource) {
211       mySourceResults.add(resolveResult);
212     }
213     else {
214       myLibResults.add(resolveResult);
215     }
216   }
217
218   @Override
219   @NotNull
220   public List<PsiElement> resultsAsList() {
221     if (!myContext.isValid()) {
222       return Collections.emptyList();
223     }
224
225     final PsiFile footholdFile = myContext.getFootholdFile();
226     if (myRelativeLevel >= 0 && footholdFile != null && !PyUserSkeletonsUtil.isUnderUserSkeletonsDirectory(footholdFile)) {
227       PsiDirectory dir = footholdFile.getContainingDirectory();
228       if (myRelativeLevel > 0) {
229         dir = ResolveImportUtil.stepBackFrom(footholdFile, myRelativeLevel);
230       }
231
232       PsiElement module = resolveModuleAt(dir);
233       if (module != null) {
234         addRoot(module, true);
235       }
236     }
237
238     final PythonPathCache cache = findMyCache();
239     final boolean mayCache = cache != null && !myWithoutRoots && !myWithoutForeign;
240     if (mayCache) {
241       final List<PsiElement> cachedResults = cache.get(myQualifiedName);
242       if (cachedResults != null) {
243         mySourceResults.addAll(cachedResults);
244         return Lists.newArrayList(mySourceResults);
245       }
246     }
247
248     if (!myWithoutRoots) {
249       addResultsFromRoots();
250     }
251     else if (footholdFile != null){
252       addResultsFromSkeletons(footholdFile);
253     }
254
255     mySourceResults.addAll(myLibResults);
256     myLibResults.clear();
257
258     if (!myWithoutForeign) {
259       for (PyImportResolver resolver : Extensions.getExtensions(PyImportResolver.EP_NAME)) {
260         PsiElement foreign = resolver.resolveImportReference(myQualifiedName, myContext);
261         if (foreign != null) {
262           myForeignResults.add(foreign);
263         }
264       }
265       mySourceResults.addAll(myForeignResults);
266       myForeignResults.clear();
267     }
268
269     final ArrayList<PsiElement> results = Lists.newArrayList(mySourceResults);
270     if (mayCache) {
271       cache.put(myQualifiedName, results);
272     }
273     return results;
274   }
275
276   /**
277    * Resolve relative imports from sdk root to the skeleton dir
278    */
279   private void addResultsFromSkeletons(@NotNull final PsiFile foothold) {
280     final boolean inSource = FileIndexFacade.getInstance(foothold.getProject()).isInContent(foothold.getVirtualFile());
281     if (inSource) return;
282     PsiDirectory containingDirectory = foothold.getContainingDirectory();
283     if (myRelativeLevel > 0) {
284       containingDirectory = ResolveImportUtil.stepBackFrom(foothold, myRelativeLevel);
285     }
286     if (containingDirectory != null) {
287       final QualifiedName containingPath = QualifiedNameFinder.findCanonicalImportPath(containingDirectory, null);
288       if (containingPath != null && containingPath.getComponentCount() > 0) {
289         final QualifiedName absolutePath = containingPath.append(myQualifiedName.toString());
290         final QualifiedNameResolverImpl absoluteVisitor =
291           (QualifiedNameResolverImpl)new QualifiedNameResolverImpl(absolutePath).fromElement(foothold);
292
293         final Sdk sdk = PythonSdkType.getSdk(foothold);
294         if (sdk == null) return;
295         final VirtualFile skeletonsDir = PySdkUtil.findSkeletonsDir(sdk);
296         if (skeletonsDir == null) return;
297         final PsiDirectory directory = myContext.getPsiManager().findDirectory(skeletonsDir);
298         final PsiElement psiElement = absoluteVisitor.resolveModuleAt(directory);
299         if (psiElement != null)
300           myLibResults.add(psiElement);
301       }
302     }
303   }
304
305   private void addResultsFromRoots() {
306     if (myVisitAllModules) {
307       for (Module mod : ModuleManager.getInstance(myContext.getProject()).getModules()) {
308         RootVisitorHost.visitRoots(mod, true, this);
309       }
310
311       if (myContext.getSdk() != null) {
312         RootVisitorHost.visitSdkRoots(myContext.getSdk(), this);
313       }
314       else if (myContext.getFootholdFile() != null) {
315         RootVisitorHost.visitSdkRoots(myContext.getFootholdFile(), this);
316       }
317     }
318     else if (myContext.getModule() != null) {
319       final boolean otherSdk = withOtherSdk();
320       RootVisitorHost.visitRoots(myContext.getModule(), otherSdk, this);
321       if (otherSdk) {
322         RootVisitorHost.visitSdkRoots(myContext.getSdk(), this);
323       }
324     }
325     else if (myContext.getFootholdFile() != null) {
326       RootVisitorHost.visitSdkRoots(myContext.getFootholdFile(), this);
327     }
328     else if (myContext.getSdk() != null) {
329       RootVisitorHost.visitSdkRoots(myContext.getSdk(), this);
330     }
331     else {
332       throw new IllegalStateException();
333     }
334   }
335
336   @Override
337   @Nullable
338   public PsiElement firstResult() {
339     final List<PsiElement> results = resultsAsList();
340     return results.size() > 0 ? results.get(0) : null;
341   }
342
343   @Override
344   @NotNull
345   public <T extends PsiElement> List<T> resultsOfType(Class<T> clazz) {
346     List<T> result = new ArrayList<T>();
347     for (PsiElement element : resultsAsList()) {
348       if (clazz.isInstance(element)) {
349         //noinspection unchecked
350         result.add((T)element);
351       }
352     }
353     return result;
354   }
355
356   @Override
357   @Nullable
358   public <T extends PsiElement> T firstResultOfType(Class<T> clazz) {
359     final List<T> list = resultsOfType(clazz);
360     return list.size() > 0 ? list.get(0) : null;
361   }
362
363   private boolean withOtherSdk() {
364     return myContext.getSdk() != null && myContext.getSdk() != PythonSdkType.findPythonSdk(myContext.getModule());
365   }
366
367   @Nullable
368   private PythonPathCache findMyCache() {
369     if (myVisitAllModules) {
370       return null;
371     }
372     if (myContext.getModule() != null) {
373       return withOtherSdk() ? null : PythonModulePathCache.getInstance(myContext.getModule());
374     }
375     if (myContext.getFootholdFile() != null) {
376       final Sdk sdk = PyBuiltinCache.findSdkForNonModuleFile(myContext.getFootholdFile());
377       if (sdk != null) {
378         return PythonSdkPathCache.getInstance(myContext.getProject(), sdk);
379       }
380     }
381     return null;
382   }
383
384   @Nullable
385   private PsiElement resolveInRoot(VirtualFile root) {
386     if (!root.isDirectory()) {
387       // if we have added a file as a root, it's unlikely that we'll be able to resolve anything under it in 'files only' resolve mode
388       return null;
389     }
390     return resolveModuleAt(myContext.getPsiManager().findDirectory(root));
391   }
392
393   /**
394    * Searches for a module at given directory, unwinding qualifiers and traversing directories as needed.
395    *
396    * @param directory where to start from; top qualifier will be searched for here.
397    */
398   @Nullable
399   public PsiElement resolveModuleAt(@Nullable PsiDirectory directory) {
400     // prerequisites
401     if (directory == null || !directory.isValid()) return null;
402
403     PsiElement seeker = directory;
404     for (String name : myQualifiedName.getComponents()) {
405       if (name == null) {
406         return null;
407       }
408       seeker = ResolveImportUtil.resolveChild(seeker, name, myContext.getFootholdFile(), !myWithMembers, myCheckForPackage);
409     }
410     return seeker;
411   }
412
413   @Override
414   public Module getModule() {
415     return myContext.getModule();
416   }
417
418   @Nullable
419   @Override
420   public <T extends PsiNamedElement> T resolveTopLevelMember(@NotNull final Class<T> aClass) {
421     Preconditions.checkState(getModule() != null, "Module is not set");
422     final String memberName = myQualifiedName.getLastComponent();
423     if (memberName == null) {
424       return null;
425     }
426     final PyFile file =
427       new QualifiedNameResolverImpl(myQualifiedName.removeLastComponent()).fromModule(getModule()).firstResultOfType(PyFile.class);
428     if (file == null) {
429       return null;
430     }
431     for (final T element : PsiTreeUtil.getChildrenOfTypeAsList(file, aClass)) {
432       if (memberName.equals(element.getName())) {
433         return element;
434       }
435     }
436     return null;
437   }
438 }