8fb2e15d81274fb2e15a83361b58b99de9754773
[idea/community.git] / python / src / com / jetbrains / python / codeInsight / intentions / PyConvertTypeCommentToVariableAnnotation.java
1 /*
2  * Copyright 2000-2016 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.openapi.editor.Document;
19 import com.intellij.openapi.editor.Editor;
20 import com.intellij.openapi.project.Project;
21 import com.intellij.psi.PsiComment;
22 import com.intellij.psi.PsiDocumentManager;
23 import com.intellij.psi.PsiElement;
24 import com.intellij.psi.PsiFile;
25 import com.intellij.psi.util.PsiTreeUtil;
26 import com.intellij.util.IncorrectOperationException;
27 import com.jetbrains.python.PyBundle;
28 import com.jetbrains.python.codeInsight.PyTypingTypeProvider;
29 import com.jetbrains.python.psi.*;
30 import com.jetbrains.python.psi.impl.PyPsiUtils;
31 import org.jetbrains.annotations.Nls;
32 import org.jetbrains.annotations.NotNull;
33 import org.jetbrains.annotations.Nullable;
34
35 public class PyConvertTypeCommentToVariableAnnotation extends PyBaseIntentionAction {
36   @Override
37   public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
38     final PsiComment comment = findCommentUnderCaret(editor, file);
39     if (comment != null) {
40       final PyTargetExpression target = findTypeCommentTarget(comment);
41       if (target != null) {
42         final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText());
43         final PsiElement prev = PyPsiUtils.getPrevNonWhitespaceSibling(comment);
44         final int commentStart = prev != null ? prev.getTextRange().getEndOffset() : comment.getTextRange().getStartOffset();
45         final int commentEnd = comment.getTextRange().getEndOffset();
46         final Document document = editor.getDocument();
47         runWithDocumentReleasedAndCommitted(project, document, () -> {
48           document.deleteString(commentStart, commentEnd);
49           document.insertString(target.getTextRange().getEndOffset(), ": " + annotation);
50         });
51       }
52     }
53   }
54
55   @NotNull
56   @Override
57   public String getText() {
58     return PyBundle.message("INTN.convert.type.comment.to.variable.annotation.text");
59   }
60
61   @Nls
62   @NotNull
63   @Override
64   public String getFamilyName() {
65     return PyBundle.message("INTN.convert.type.comment.to.variable.annotation.family");
66   }
67
68   @Override
69   public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
70     if (file instanceof PyFile && LanguageLevel.forElement(file).isAtLeast(LanguageLevel.PYTHON36)) {
71       final PsiComment comment = findCommentUnderCaret(editor, file);
72       return comment != null && isSuitableTypeComment(comment);
73     }
74     return false;
75   }
76
77   @Nullable
78   private static PsiComment findCommentUnderCaret(@NotNull Editor editor, @NotNull PsiFile file) {
79     final PsiElement element = file.findElementAt(editor.getCaretModel().getOffset());
80     return PsiTreeUtil.getParentOfType(element, PsiComment.class, false);
81   }
82
83   private boolean isSuitableTypeComment(@NotNull PsiComment comment) {
84     final String annotation = PyTypingTypeProvider.getTypeCommentValue(comment.getText());
85     return annotation != null && findTypeCommentTarget(comment) != null;
86   }
87
88   @Nullable
89   private PyTargetExpression findTypeCommentTarget(@NotNull PsiComment comment) {
90     final PsiElement parent = comment.getParent();
91     if (parent instanceof PyAssignmentStatement) {
92       final PyAssignmentStatement assignment = (PyAssignmentStatement)parent;
93       final PyExpression[] rawTargets = assignment.getRawTargets();
94       if (rawTargets.length == 1 && rawTargets[0] instanceof PyTargetExpression) {
95         final PyTargetExpression target = (PyTargetExpression)rawTargets[0];
96         if (target.getTypeComment() == comment) {
97           return target;
98         }
99       }
100     }
101     return null;
102   }
103
104   public static void runWithDocumentReleasedAndCommitted(@NotNull Project project, @NotNull Document document, @NotNull Runnable runnable) {
105     final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
106     manager.doPostponedOperationsAndUnblockDocument(document);
107     try {
108       runnable.run();
109     }
110     finally {
111       manager.commitDocument(document);
112     }
113   }
114 }