python/src/com/jetbrains/python/codeInsight/intentions/PyConvertLocalFunctionToTopLevelFunction.java

   1 /*
   2  * Copyright 2000-2015 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.intentions;
  17
  18 import com.intellij.codeInsight.controlflow.ControlFlow;
  19 import com.intellij.codeInsight.controlflow.Instruction;
  20 import com.intellij.codeInsight.intention.impl.BaseIntentionAction;
  21 import com.intellij.openapi.editor.Editor;
  22 import com.intellij.openapi.project.Project;
  23 import com.intellij.openapi.roots.ProjectRootManager;
  24 import com.intellij.openapi.ui.MessageType;
  25 import com.intellij.openapi.ui.popup.Balloon;
  26 import com.intellij.openapi.ui.popup.JBPopupFactory;
  27 import com.intellij.openapi.util.text.StringUtil;
  28 import com.intellij.openapi.vfs.VirtualFile;
  29 import com.intellij.psi.PsiElement;
  30 import com.intellij.psi.PsiFile;
  31 import com.intellij.psi.util.PsiTreeUtil;
  32 import com.intellij.usageView.UsageInfo;
  33 import com.intellij.util.IncorrectOperationException;
  34 import com.intellij.util.containers.ContainerUtil;
  35 import com.jetbrains.python.PyBundle;
  36 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
  37 import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
  38 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
  39 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
  40 import com.jetbrains.python.psi.*;
  41 import com.jetbrains.python.psi.impl.PyPsiUtils;
  42 import com.jetbrains.python.psi.resolve.PyResolveContext;
  43 import com.jetbrains.python.psi.types.TypeEvalContext;
  44 import com.jetbrains.python.refactoring.PyRefactoringUtil;
  45 import org.jetbrains.annotations.Nls;
  46 import org.jetbrains.annotations.NotNull;
  47 import org.jetbrains.annotations.Nullable;
  48
  49 import java.util.*;
  50
  51 import static com.jetbrains.python.psi.PyUtil.as;
  52
  53 /**
  54  * @author Mikhail Golubev
  55  */
  56 public class PyConvertLocalFunctionToTopLevelFunction extends BaseIntentionAction {
  57   public PyConvertLocalFunctionToTopLevelFunction() {
  58     setText(PyBundle.message("INTN.convert.local.function.to.top.level.function"));
  59   }
  60
  61   @Nls
  62   @NotNull
  63   @Override
  64   public String getFamilyName() {
  65     return PyBundle.message("INTN.convert.local.function.to.top.level.function");
  66   }
  67
  68   @Override
  69   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
  70     final PyFunction nestedFunction = findNestedFunctionUnderCaret(editor, file);
  71     return nestedFunction != null;
  72   }
  73
  74   @Nullable
  75   private static PyFunction findNestedFunctionUnderCaret(Editor editor, PsiFile file) {
  76     if (!(file instanceof PyFile)) return null;
  77     final PsiElement element = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
  78     if (element == null) {
  79       return null;
  80     }
  81     PyFunction result = null;
  82     if (isLocalFunction(element.getParent()) && ((PyFunction)element.getParent()).getNameIdentifier() == element) {
  83       result = (PyFunction)element.getParent();
  84     }
  85     else {
  86       final PyReferenceExpression refExpr = PsiTreeUtil.getParentOfType(element, PyReferenceExpression.class);
  87       if (refExpr == null) {
  88         return null;
  89       }
  90       final PsiElement resolved = refExpr.getReference().resolve();
  91       if (isLocalFunction(resolved)) {
  92         result = (PyFunction)resolved;
  93       }
  94     }
  95     if (result != null) {
  96       final VirtualFile virtualFile = result.getContainingFile().getVirtualFile();
  97       if (virtualFile != null && ProjectRootManager.getInstance(file.getProject()).getFileIndex().isInLibraryClasses(virtualFile)) {
  98         return null;
  99       }
 100     }
 101     return result;
 102   }
 103
 104   private static boolean isLocalFunction(@Nullable PsiElement resolved) {
 105     if (resolved instanceof PyFunction && PsiTreeUtil.getParentOfType(resolved, ScopeOwner.class, true) instanceof PyFunction) {
 106       return true;
 107     }
 108     return false;
 109   }
 110
 111   @Override
 112   public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
 113     final PyResolveContext context = PyResolveContext.defaultContext().withTypeEvalContext(TypeEvalContext.userInitiated(project, file));
 114     final PyFunction function = findNestedFunctionUnderCaret(editor, file);
 115     assert function != null;
 116     final Set<String> enclosingScopeReads = new LinkedHashSet<String>();
 117     final Collection<ScopeOwner> scopeOwners = PsiTreeUtil.collectElementsOfType(function, ScopeOwner.class);
 118     for (ScopeOwner owner : scopeOwners) {
 119       final AnalysisResult scope = findReadsFromEnclosingScope(owner, function, context);
 120       if (!scope.nonlocalWritesToEnclosingScope.isEmpty()) {
 121         showErrorBalloon(editor, PyBundle.message("INTN.convert.local.function.to.top.level.function.nonlocal"));
 122         return;
 123       }
 124       for (PsiElement element : scope.readFromEnclosingScope) {
 125         if (element instanceof PyElement) {
 126           ContainerUtil.addIfNotNull(enclosingScopeReads, ((PyElement)element).getName());
 127         }
 128       }
 129     }
 130     final String commaSeparatedNames = StringUtil.join(enclosingScopeReads, ", ");
 131
 132     // Update existing usages
 133     final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
 134     for (UsageInfo usage : PyRefactoringUtil.findUsages(function, false)) {
 135       final PsiElement element = usage.getElement();
 136       if (element != null) {
 137         final PyCallExpression parentCall = as(element.getParent(), PyCallExpression.class);
 138         if (parentCall != null) {
 139           final PyArgumentList argList = parentCall.getArgumentList();
 140           if (argList != null) {
 141             final StringBuilder argListText = new StringBuilder(argList.getText());
 142             argListText.insert(1, commaSeparatedNames + (argList.getArguments().length > 0 ? ", " : ""));
 143             argList.replace(elementGenerator.createArgumentList(LanguageLevel.forElement(element), argListText.toString()));
 144           }
 145         }
 146       }
 147     }
 148
 149     // Replace function
 150     PyFunction copiedFunction = (PyFunction)function.copy();
 151     final PyParameterList paramList = copiedFunction.getParameterList();
 152     final StringBuilder paramListText = new StringBuilder(paramList.getText());
 153     paramListText.insert(1, commaSeparatedNames + (paramList.getParameters().length > 0 ? ", " : ""));
 154     paramList.replace(elementGenerator.createParameterList(LanguageLevel.forElement(function), paramListText.toString()));
 155
 156     // See AddImportHelper.getFileInsertPosition()
 157     final PsiElement anchor = PyPsiUtils.getParentRightBefore(function, file);
 158
 159     copiedFunction = (PyFunction)file.addAfter(copiedFunction, anchor);
 160     function.delete();
 161
 162     editor.getSelectionModel().removeSelection();
 163     editor.getCaretModel().moveToOffset(copiedFunction.getTextOffset());
 164   }
 165
 166   @NotNull
 167   private static AnalysisResult findReadsFromEnclosingScope(@NotNull ScopeOwner owner,
 168                                                             @NotNull PyFunction targetFunction,
 169                                                             @NotNull PyResolveContext context) {
 170     final ControlFlow controlFlow = ControlFlowCache.getControlFlow(owner);
 171     final List<PsiElement> readFromEnclosingScope = new ArrayList<PsiElement>();
 172     final List<PyTargetExpression> nonlocalWrites = new ArrayList<PyTargetExpression>();
 173     for (Instruction instruction : controlFlow.getInstructions()) {
 174       if (instruction instanceof ReadWriteInstruction) {
 175         final ReadWriteInstruction readWriteInstruction = (ReadWriteInstruction)instruction;
 176         final PsiElement element = readWriteInstruction.getElement();
 177         if (element == null) {
 178           continue;
 179         }
 180         if (readWriteInstruction.getAccess().isReadAccess()) {
 181           for (PsiElement resolved : PyUtil.multiResolveTopPriority(element, context)) {
 182             if (resolved != null && isFromEnclosingScope(resolved, targetFunction)) {
 183               readFromEnclosingScope.add(element);
 184               break;
 185             }
 186           }
 187         }
 188         if (readWriteInstruction.getAccess().isWriteAccess()) {
 189           if (element instanceof PyTargetExpression && element.getParent() instanceof PyNonlocalStatement) {
 190             for (PsiElement resolved : PyUtil.multiResolveTopPriority(element, context)) {
 191               if (resolved != null && isFromEnclosingScope(resolved, targetFunction)) {
 192                 nonlocalWrites.add((PyTargetExpression)element);
 193                 break;
 194               }
 195             }
 196           }
 197         }
 198       }
 199     }
 200     return new AnalysisResult(readFromEnclosingScope, nonlocalWrites);
 201   }
 202
 203   private static class AnalysisResult {
 204     final List<PsiElement> readFromEnclosingScope;
 205     final List<PyTargetExpression> nonlocalWritesToEnclosingScope;
 206
 207     public AnalysisResult(@NotNull List<PsiElement> readFromEnclosingScope, @NotNull List<PyTargetExpression> nonlocalWrites) {
 208       this.readFromEnclosingScope = readFromEnclosingScope;
 209       this.nonlocalWritesToEnclosingScope = nonlocalWrites;
 210     }
 211   }
 212
 213   private static boolean isFromEnclosingScope(@NotNull PsiElement element, @NotNull PyFunction targetFunction) {
 214     return !PsiTreeUtil.isAncestor(targetFunction, element, false) && !(ScopeUtil.getScopeOwner(element) instanceof PsiFile);
 215   }
 216
 217   private static void showErrorBalloon(@NotNull Editor editor, @NotNull String message) {
 218     final JBPopupFactory popupFactory = JBPopupFactory.getInstance();
 219     popupFactory.createHtmlTextBalloonBuilder(message, MessageType.ERROR, null)
 220                 .setDisposable(editor.getProject())
 221                 .createBalloon()
 222                 .show(popupFactory.guessBestPopupLocation(editor), Balloon.Position.below);
 223   }
 224 }