| /* |
| * Copyright 2000-2013 JetBrains s.r.o. |
| * |
| * Licensed under the Apache License, Version 2.0 (the "License"); |
| * you may not use this file except in compliance with the License. |
| * You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package com.jetbrains.python.codeInsight.completion; |
| |
| import com.intellij.codeInsight.completion.*; |
| import com.intellij.codeInsight.lookup.LookupElement; |
| import com.intellij.codeInsight.lookup.LookupElementBuilder; |
| import com.intellij.navigation.NavigationItem; |
| import com.intellij.openapi.application.Result; |
| import com.intellij.openapi.command.WriteCommandAction; |
| import com.intellij.openapi.editor.Document; |
| import com.intellij.openapi.project.Project; |
| import com.intellij.openapi.util.Condition; |
| import com.intellij.openapi.util.Conditions; |
| import com.intellij.openapi.vfs.VirtualFile; |
| import com.intellij.psi.*; |
| import com.intellij.psi.search.FileTypeIndex; |
| import com.intellij.psi.search.GlobalSearchScope; |
| import com.intellij.psi.stubs.StubIndex; |
| import com.intellij.psi.stubs.StubIndexKey; |
| import com.intellij.psi.util.PsiTreeUtil; |
| import com.jetbrains.python.PythonFileType; |
| import com.jetbrains.python.codeInsight.imports.AddImportHelper; |
| import com.jetbrains.python.codeInsight.imports.PythonReferenceImporter; |
| import com.jetbrains.python.psi.*; |
| import com.jetbrains.python.psi.search.PyProjectScopeBuilder; |
| import com.jetbrains.python.psi.stubs.PyClassNameIndex; |
| import com.jetbrains.python.psi.stubs.PyFunctionNameIndex; |
| import com.jetbrains.python.psi.stubs.PyVariableNameIndex; |
| import com.jetbrains.python.psi.types.PyModuleType; |
| import org.jetbrains.annotations.NotNull; |
| |
| import java.util.Collection; |
| |
| /** |
| * @author yole |
| */ |
| public class PyClassNameCompletionContributor extends CompletionContributor { |
| |
| @Override |
| public void fillCompletionVariants(@NotNull CompletionParameters parameters, @NotNull CompletionResultSet result) { |
| if (parameters.isExtendedCompletion()) { |
| final PsiElement element = parameters.getPosition(); |
| final PsiElement parent = element.getParent(); |
| if (parent instanceof PyReferenceExpression && ((PyReferenceExpression)parent).isQualified()) { |
| return; |
| } |
| if (parent instanceof PyStringLiteralExpression) { |
| String prefix = parent.getText().substring(0, parameters.getOffset() - parent.getTextRange().getStartOffset()); |
| if (prefix.contains(".")) { |
| return; |
| } |
| } |
| final FileViewProvider provider = element.getContainingFile().getViewProvider(); |
| if (provider instanceof MultiplePsiFilesPerDocumentFileViewProvider) return; |
| if (PsiTreeUtil.getParentOfType(element, PyImportStatementBase.class) != null) { |
| return; |
| } |
| final PsiFile originalFile = parameters.getOriginalFile(); |
| addVariantsFromIndex(result, originalFile, PyClassNameIndex.KEY, |
| parent instanceof PyStringLiteralExpression ? STRING_LITERAL_INSERT_HANDLER : IMPORTING_INSERT_HANDLER, |
| Conditions.<PyClass>alwaysTrue(), PyClass.class); |
| addVariantsFromIndex(result, originalFile, PyFunctionNameIndex.KEY, |
| getFunctionInsertHandler(parent), IS_TOPLEVEL, PyFunction.class); |
| addVariantsFromIndex(result, originalFile, PyVariableNameIndex.KEY, |
| parent instanceof PyStringLiteralExpression ? STRING_LITERAL_INSERT_HANDLER : IMPORTING_INSERT_HANDLER, |
| IS_TOPLEVEL, PyTargetExpression.class); |
| addVariantsFromModules(result, originalFile, parent instanceof PyStringLiteralExpression); |
| } |
| } |
| |
| private static InsertHandler<LookupElement> getFunctionInsertHandler(PsiElement parent) { |
| if (parent instanceof PyStringLiteralExpression) { |
| return STRING_LITERAL_INSERT_HANDLER; |
| } |
| if (parent.getParent() instanceof PyDecorator) { |
| return IMPORTING_INSERT_HANDLER; |
| } |
| return FUNCTION_INSERT_HANDLER; |
| } |
| |
| private static void addVariantsFromModules(CompletionResultSet result, PsiFile targetFile, boolean inStringLiteral) { |
| Collection<VirtualFile> files = FileTypeIndex.getFiles(PythonFileType.INSTANCE, PyProjectScopeBuilder.excludeSdkTestsScope(targetFile)); |
| for (VirtualFile file : files) { |
| PsiFile pyFile = targetFile.getManager().findFile(file); |
| PsiFileSystemItem importable = (PsiFileSystemItem) PyUtil.turnInitIntoDir(pyFile); |
| if (importable == null) continue; |
| if (PythonReferenceImporter.isImportableModule(targetFile, importable)) { |
| LookupElementBuilder element = PyModuleType.buildFileLookupElement(importable, null); |
| if (element != null) { |
| result.addElement(element.withInsertHandler(inStringLiteral ? STRING_LITERAL_INSERT_HANDLER : IMPORTING_INSERT_HANDLER)); |
| } |
| } |
| } |
| } |
| |
| private static Condition<PsiElement> IS_TOPLEVEL = new Condition<PsiElement>() { |
| @Override |
| public boolean value(PsiElement element) { |
| return PyUtil.isTopLevel(element); |
| } |
| }; |
| |
| private static <T extends PsiNamedElement> void addVariantsFromIndex(final CompletionResultSet resultSet, |
| final PsiFile targetFile, |
| final StubIndexKey<String, T> key, |
| final InsertHandler<LookupElement> insertHandler, |
| final Condition<? super T> condition, Class<T> elementClass) { |
| final Project project = targetFile.getProject(); |
| GlobalSearchScope scope = PyProjectScopeBuilder.excludeSdkTestsScope(targetFile); |
| |
| Collection<String> keys = StubIndex.getInstance().getAllKeys(key, project); |
| for (final String elementName : CompletionUtil.sortMatching(resultSet.getPrefixMatcher(), keys)) { |
| for (T element : StubIndex.getElements(key, elementName, project, scope, elementClass)) { |
| if (condition.value(element)) { |
| resultSet.addElement(LookupElementBuilder.createWithIcon(element) |
| .withTailText(" " + ((NavigationItem)element).getPresentation().getLocationString(), true) |
| .withInsertHandler(insertHandler)); |
| } |
| } |
| } |
| } |
| |
| private static final InsertHandler<LookupElement> IMPORTING_INSERT_HANDLER = new InsertHandler<LookupElement>() { |
| public void handleInsert(final InsertionContext context, final LookupElement item) { |
| addImportForLookupElement(context, item, context.getTailOffset() - 1); |
| } |
| }; |
| |
| |
| private static final InsertHandler<LookupElement> FUNCTION_INSERT_HANDLER = new PyFunctionInsertHandler() { |
| public void handleInsert(final InsertionContext context, final LookupElement item) { |
| int tailOffset = context.getTailOffset()-1; |
| super.handleInsert(context, item); // adds parentheses, modifies tail offset |
| context.commitDocument(); |
| addImportForLookupElement(context, item, tailOffset); |
| } |
| }; |
| |
| private static final InsertHandler<LookupElement> STRING_LITERAL_INSERT_HANDLER = new InsertHandler<LookupElement>() { |
| @Override |
| public void handleInsert(InsertionContext context, LookupElement item) { |
| PsiElement element = item.getPsiElement(); |
| if (element instanceof PyQualifiedNameOwner) { |
| String qName = ((PyQualifiedNameOwner) element).getQualifiedName(); |
| String name = ((PyQualifiedNameOwner) element).getName(); |
| if (qName != null && name != null) { |
| String qNamePrefix = qName.substring(0, qName.length()-name.length()); |
| context.getDocument().insertString(context.getStartOffset(), qNamePrefix); |
| } |
| } |
| } |
| }; |
| |
| private static void addImportForLookupElement(final InsertionContext context, final LookupElement item, final int tailOffset) { |
| PsiDocumentManager manager = PsiDocumentManager.getInstance(context.getProject()); |
| Document document = manager.getDocument(context.getFile()); |
| if (document != null) { |
| manager.commitDocument(document); |
| } |
| final PsiReference ref = context.getFile().findReferenceAt(tailOffset); |
| if (ref == null || ref.resolve() == item.getObject()) { |
| // no import statement needed |
| return; |
| } |
| new WriteCommandAction(context.getProject(), context.getFile()) { |
| @Override |
| protected void run(Result result) throws Throwable { |
| AddImportHelper.addImport((PsiNamedElement)item.getObject(), context.getFile(), (PyElement)ref.getElement()); |
| } |
| }.execute(); |
| } |
| } |