diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java index 075bd5f4c841..dd2b90d0256d 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java @@ -16,13 +16,18 @@ package com.jetbrains.python.codeInsight.intentions; import com.google.common.base.Function; +import com.google.common.collect.Lists; import com.intellij.codeInsight.CodeInsightUtilCore; import com.intellij.codeInsight.intention.IntentionAction; import com.intellij.codeInsight.template.*; +import com.intellij.openapi.editor.Document; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.fileEditor.FileEditorManager; import com.intellij.openapi.fileEditor.OpenFileDescriptor; import com.intellij.openapi.project.Project; +import com.intellij.openapi.util.Pair; +import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiDocumentManager; import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; import com.intellij.util.IncorrectOperationException; @@ -31,8 +36,9 @@ import com.jetbrains.python.documentation.doctest.PyDocstringFile; import com.jetbrains.python.psi.*; import org.jetbrains.annotations.NotNull; -import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateParameter; -import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.annotateReturnType; +import java.util.List; + +import static com.jetbrains.python.codeInsight.intentions.SpecifyTypeInPy3AnnotationsIntention.*; import static com.jetbrains.python.codeInsight.intentions.TypeIntention.getCallable; import static com.jetbrains.python.codeInsight.intentions.TypeIntention.resolvesToFunction; @@ -55,8 +61,6 @@ public class PyAnnotateTypesIntention implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false; - if (!LanguageLevel.forElement(file).isPy3K()) return false; - updateText(); final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); @@ -78,7 +82,106 @@ public class PyAnnotateTypesIntention implements IntentionAction { final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); final PyCallable callable = getCallable(elementAt); + if (isPy3k(file)) { + generatePy3kTypeAnnotations(project, editor, elementAt, callable); + } + else { + if (callable instanceof PyFunction) { + generateTypeCommentAnnotations(project, editor, elementAt, (PyFunction)callable); + } + } + } + private static void generateTypeCommentAnnotations(Project project, Editor editor, PsiElement at, PyFunction function) { + + StringBuilder replacementTextBuilder = new StringBuilder("# type: ("); + + PyParameter[] params = function.getParameterList().getParameters(); + + List> templates = Lists.newArrayList(); + + for (int i = 0; i < params.length; i++) { + String type = parameterType(params[i]); + + templates.add(Pair.create(replacementTextBuilder.length(), type)); + + replacementTextBuilder.append(type); + + if (i < params.length - 1) { + replacementTextBuilder.append(", "); + } + } + + replacementTextBuilder.append(") -> "); + + String returnType = returnType(function); + templates.add(Pair.create(replacementTextBuilder.length(), returnType)); + + replacementTextBuilder.append(returnType); + + final PyStatementList statements = function.getStatementList(); + final String indentation = PyIndentUtil.getExpectedElementIndent(statements); + replacementTextBuilder.insert(0, indentation); + replacementTextBuilder.insert(0, "\n"); + + + final PsiDocumentManager manager = PsiDocumentManager.getInstance(project); + final Document document = manager.getDocument(function.getContainingFile()); + + if (document != null) { + final PsiElement beforeStatements = statements.getPrevSibling(); + int offset = beforeStatements.getTextRange().getStartOffset(); + if (":".equals(beforeStatements.getText())) { + offset += 1; + } + try { + document.insertString(offset, replacementTextBuilder.toString()); + } + finally { + manager.commitDocument(document); + } + + + function = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function); + + if (function != null) { + final TemplateBuilder builder = + TemplateBuilderFactory.getInstance().createTemplateBuilder(function); + + for (Pair template : templates) { + builder.replaceRange(TextRange.from( + offset - function.getTextRange().getStartOffset() + replacementTextBuilder.toString().indexOf('#') + template.first, + template.second.length()), template.second); + } + + startTemplate(project, function, builder); + } + } + } + + private static void startTemplate(Project project, PyCallable callable, TemplateBuilder builder) { + final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + ; + + int offset = callable.getTextRange().getStartOffset(); + + final OpenFileDescriptor descriptor = new OpenFileDescriptor( + project, + callable.getContainingFile().getVirtualFile(), + offset + ); + final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + if (targetEditor != null) { + targetEditor.getCaretModel().moveToOffset(offset); + TemplateManager.getInstance(project).startTemplate(targetEditor, template); + } + } + + private static boolean isPy3k(PsiFile file) { + return LanguageLevel.forElement(file).isPy3K(); + } + + private static void generatePy3kTypeAnnotations(@NotNull Project project, Editor editor, PsiElement elementAt, PyCallable callable) { final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(callable); PyExpression returnType = annotateReturnType(project, editor.getDocument(), elementAt, false); @@ -88,7 +191,7 @@ public class PyAnnotateTypesIntention implements IntentionAction { } if (callable instanceof PyFunction) { - PyFunction function = (PyFunction) callable; + PyFunction function = (PyFunction)callable; PyParameter[] params = function.getParameterList().getParameters(); for (int i = params.length - 1; i >= 0; i--) { @@ -112,20 +215,7 @@ public class PyAnnotateTypesIntention implements IntentionAction { } } if (callable != null) { - final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - - int offset = callable.getTextRange().getStartOffset(); - - final OpenFileDescriptor descriptor = new OpenFileDescriptor( - project, - callable.getContainingFile().getVirtualFile(), - offset - ); - final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); - if (targetEditor != null) { - targetEditor.getCaretModel().moveToOffset(offset); - TemplateManager.getInstance(project).startTemplate(targetEditor, template); - } + startTemplate(project, callable, builder); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index 4afc6751a00d..180a2165d9c4 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -91,16 +91,9 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { final String defaultParamText = defaultParamValue == null ? null : defaultParamValue.getText(); - String paramType = PyNames.OBJECT; + String paramType = parameterType(parameter); - PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class); - if (function != null) { - final PySignature signature = PySignatureCacheManager.getInstance(parameter.getProject()).findSignature( - function); - if (signature != null) { - paramType = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(paramName), paramType); - } - } + final PyNamedParameter namedParameter = elementGenerator.createParameter(paramName, defaultParamText, paramType, LanguageLevel.forElement(parameter)); @@ -124,18 +117,39 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { return parameter; } + static String parameterType(PyParameter parameter) { + String paramType = PyNames.OBJECT; + + PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class); + if (function != null) { + final PySignature signature = PySignatureCacheManager.getInstance(parameter.getProject()).findSignature( + function); + String parameterName = parameter.getName(); + if (signature != null && parameterName != null) { + paramType = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(parameterName), paramType); + } + } + return paramType; + } + + + static String returnType(@NotNull PyFunction function) { + String returnType = PyNames.OBJECT; + final PySignature signature = PySignatureCacheManager.getInstance(function.getProject()).findSignature( + function); + if (signature != null) { + returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType); + } + return returnType; + } + public static PyExpression annotateReturnType(Project project, Document document, PsiElement resolved, boolean createTemplate) { PyCallable callable = getCallable(resolved); - String returnType = PyNames.OBJECT; - if (callable instanceof PyFunction) { PyFunction function = (PyFunction)callable; - final PySignature signature = PySignatureCacheManager.getInstance(project).findSignature( - function); - if (signature != null) { - returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType); - } + + String returnType = returnType(function); final String annotationText = " -> " + returnType; @@ -143,20 +157,22 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { assert prevElem != null; final PsiDocumentManager manager = PsiDocumentManager.getInstance(project); - Document documentWithCollable = manager.getDocument(callable.getContainingFile()); - try { - final TextRange range = prevElem.getTextRange(); - manager.doPostponedOperationsAndUnblockDocument(documentWithCollable); - if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) { - documentWithCollable.insertString(range.getStartOffset(), annotationText); + Document documentWithCallable = manager.getDocument(callable.getContainingFile()); + if (documentWithCallable != null) { + try { + final TextRange range = prevElem.getTextRange(); + manager.doPostponedOperationsAndUnblockDocument(documentWithCallable); + if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) { + documentWithCallable.insertString(range.getStartOffset(), annotationText); + } + else { + documentWithCallable.insertString(range.getEndOffset(), annotationText + ":"); + } } - else { - documentWithCollable.insertString(range.getEndOffset(), annotationText + ":"); + finally { + manager.commitDocument(documentWithCallable); } } - finally { - manager.commitDocument(documentWithCollable); - } callable = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(callable); diff --git a/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment.py b/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment.py new file mode 100644 index 000000000000..cb92a3843add --- /dev/null +++ b/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment.py @@ -0,0 +1,2 @@ +def foo(x, y): + pass \ No newline at end of file diff --git a/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment_after.py b/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment_after.py new file mode 100644 index 000000000000..cb635649431a --- /dev/null +++ b/python/testData/intentions/PyAnnotateTypesIntentionTest/typeComment_after.py @@ -0,0 +1,3 @@ +def foo(x, y): + # type: (object, object) -> object + pass \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java index 3d2f2d2ea0e4..77cd98c1da43 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java @@ -42,6 +42,11 @@ public class PyAnnotateTypesIntentionTest extends PyIntentionTestCase { } } + public void testTypeComment() { + doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON27); + } + + private void doTest() { doTest(PyBundle.message("INTN.annotate.types"), LanguageLevel.PYTHON30); }