diff --git a/python/src/com/jetbrains/python/PyBundle.properties b/python/src/com/jetbrains/python/PyBundle.properties index a747d89160b7..62b0f1414d56 100644 --- a/python/src/com/jetbrains/python/PyBundle.properties +++ b/python/src/com/jetbrains/python/PyBundle.properties @@ -184,9 +184,11 @@ INTN.doc.string.stub=Insert documentation string stub #SpecifyTypeInDocstringIntention INTN.specify.type=Specify type for reference in docstring +INTN.specify.return.type=Specify return type in docstring #SpecifyTypeInPy3AnnotationsIntention INTN.specify.type.in.annotation=Specify type for reference using annotation +INTN.specify.returt.type.in.annotation=Specify return type using annotation #TypeAssertionIntention INTN.insert.assertion=Insert type assertion diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java index bfd60427b2c5..42c7231c08ec 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -3,17 +3,20 @@ package com.jetbrains.python.codeInsight.intentions; import com.intellij.codeInsight.CodeInsightUtilBase; import com.intellij.codeInsight.intention.IntentionAction; import com.intellij.codeInsight.template.*; +import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Document; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.project.Project; import com.intellij.openapi.util.Pair; import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; import com.intellij.psi.PsiReference; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PyBundle; import com.jetbrains.python.PyNames; +import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.documentation.PyDocumentationSettings; import com.jetbrains.python.documentation.PythonDocumentationProvider; import com.jetbrains.python.psi.*; @@ -28,13 +31,13 @@ import org.jetbrains.annotations.NotNull; * Helps to specify type */ public class SpecifyTypeInDocstringIntention implements IntentionAction { - + private String myText = PyBundle.message("INTN.specify.type"); public SpecifyTypeInDocstringIntention() { } @NotNull public String getText() { - return PyBundle.message("INTN.specify.type"); + return myText; } @NotNull @@ -43,6 +46,19 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction { } public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { + PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) + elementAt = file.findElementAt(editor.getCaretModel().getOffset()); + + PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + if (parentFunction != null) { + final ASTNode nameNode = parentFunction.getNameNode(); + if (nameNode != null && nameNode.getPsi() == elementAt) { + myText = PyBundle.message("INTN.specify.return.type"); + return true; + } + } + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); if (problemElement == null) return false; @@ -82,10 +98,33 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction { } public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { + PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) + elementAt = file.findElementAt(editor.getCaretModel().getOffset()); + + String type = "type"; + String name = ""; + PyFunction pyFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + if (pyFunction != null) { + final ASTNode nameNode = pyFunction.getNameNode(); + if (nameNode != null && nameNode.getPsi() == elementAt) { + type = "rtype"; + } + } + + PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project); + String prefix = ":"; + if (documentationSettings.isEpydocFormat(file)) { + prefix = "@"; + } + + PsiReference reference = null; PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); if (problemElement != null) { - String name = problemElement.getName(); - PsiReference reference = problemElement.getReference(); + name = problemElement.getName(); + + reference = problemElement.getReference(); if (problemElement instanceof PyQualifiedExpression) { final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); if (qualifier != null) { @@ -93,74 +132,68 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction { name = qualifier.getText(); } } - PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project); - String prefix = ":"; - if (documentationSettings.isEpydocFormat(file)) { - prefix = "@"; - } - PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class); + } - PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class); + final ASTNode nameNode = pyFunction.getNameNode(); + if ((pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) || + elementAt == nameNode.getPsi()) { + PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression(); + int startOffset; + int endOffset; + final Pair replacementToOffset = + PythonDocumentationProvider.addParamToDocstring(pyFunction, type, name, prefix); - if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) { - PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression(); - int startOffset; - int endOffset; - final Pair replacementToOffset = - PythonDocumentationProvider.addParamToDocstring(pyFunction, "type", name, prefix); - - if (docStringExpression != null) { - final String typePattern = "type " + name + ":"; - final int index = docStringExpression.getText().indexOf(typePattern); - if (index == -1) { - PyExpression str = elementGenerator.createDocstring(replacementToOffset.getFirst()).getExpression(); - docStringExpression.replace(str); - startOffset = replacementToOffset.getSecond(); - endOffset = startOffset; - } - else { - startOffset = index + typePattern.length() + 1; - endOffset = docStringExpression.getText().indexOf("\n", startOffset); - if (endOffset == -1) endOffset = startOffset; - } - pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction); - docStringExpression = pyFunction.getDocStringExpression(); - } - else { - - final PyStatementList list = pyFunction.getStatementList(); - final Document document = editor.getDocument(); + if (docStringExpression != null) { + final String typePattern = type + " " + name + ":"; + final int index = docStringExpression.getText().indexOf(typePattern); + if (index == -1) { + PyExpression str = elementGenerator.createDocstring(replacementToOffset.getFirst()).getExpression(); + docStringExpression.replace(str); startOffset = replacementToOffset.getSecond(); - - if (list != null && list.getStatements().length != 0) { - if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(pyFunction.getTextOffset())) { - PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(pyFunction), - PyFunction.class, "def " + pyFunction.getName() + pyFunction.getParameterList().getText() - +":\n\t"+replacementToOffset.getFirst() + "\n\t" + list.getText()); - - pyFunction = (PyFunction)pyFunction.replace(func); - startOffset = replacementToOffset.getSecond() + 2; - } - else { - PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst()); - list.addBefore(str, list.getStatements()[0]); - } - } - - pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction); - docStringExpression = pyFunction.getDocStringExpression(); - endOffset = startOffset; } - assert docStringExpression != null; - int textOffSet = docStringExpression.getTextOffset(); - editor.getCaretModel().moveToOffset(textOffSet); - - final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression); - builder.replaceRange(TextRange.create(startOffset, endOffset), PyNames.OBJECT); - Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - TemplateManager.getInstance(project).startTemplate(editor, template); + else { + startOffset = index + typePattern.length() + 1; + endOffset = docStringExpression.getText().indexOf("\n", startOffset); + if (endOffset == -1) endOffset = startOffset; + } + pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction); + docStringExpression = pyFunction.getDocStringExpression(); } + else { + final PyStatementList list = pyFunction.getStatementList(); + final Document document = editor.getDocument(); + startOffset = replacementToOffset.getSecond(); + + if (list != null && list.getStatements().length != 0) { + if (document.getLineNumber(list.getTextOffset()) == document.getLineNumber(pyFunction.getTextOffset())) { + PyFunction func = elementGenerator.createFromText(LanguageLevel.forElement(pyFunction), + PyFunction.class, "def " + pyFunction.getName() + pyFunction.getParameterList().getText() + +":\n\t"+replacementToOffset.getFirst() + "\n\t" + list.getText()); + + pyFunction = (PyFunction)pyFunction.replace(func); + startOffset = replacementToOffset.getSecond() + 2; + } + else { + PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst()); + list.addBefore(str, list.getStatements()[0]); + } + } + + pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction); + docStringExpression = pyFunction.getDocStringExpression(); + + endOffset = startOffset; + } + assert docStringExpression != null; + int textOffSet = docStringExpression.getTextOffset(); + editor.getCaretModel().moveToOffset(textOffSet); + + final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(docStringExpression); + builder.replaceRange(TextRange.create(startOffset, endOffset), PyNames.OBJECT); + Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + TemplateManager.getInstance(project).startTemplate(editor, template); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index e1e098d08b06..4f91fe89dff8 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -3,6 +3,7 @@ package com.jetbrains.python.codeInsight.intentions; import com.intellij.codeInsight.CodeInsightUtilBase; import com.intellij.codeInsight.intention.IntentionAction; import com.intellij.codeInsight.template.*; +import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.fileEditor.FileEditorManager; import com.intellij.openapi.fileEditor.OpenFileDescriptor; @@ -15,6 +16,7 @@ import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.IncorrectOperationException; import com.jetbrains.python.PyBundle; import com.jetbrains.python.PyNames; +import com.jetbrains.python.PyTokenTypes; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.PyResolveContext; import com.jetbrains.python.psi.types.PyReturnTypeReference; @@ -28,13 +30,13 @@ import org.jetbrains.annotations.NotNull; * Helps to specify type in annotations in python3 */ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { - + private String myText = PyBundle.message("INTN.specify.type.in.annotation"); public SpecifyTypeInPy3AnnotationsIntention() { } @NotNull public String getText() { - return PyBundle.message("INTN.specify.type.in.annotation"); + return myText; } @NotNull @@ -44,6 +46,18 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { if (!LanguageLevel.forElement(file).isPy3K()) return false; + PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) + elementAt = file.findElementAt(editor.getCaretModel().getOffset()); + + PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + if (parentFunction != null) { + final ASTNode nameNode = parentFunction.getNameNode(); + if (nameNode != null && nameNode.getPsi() == elementAt) { + myText = PyBundle.message("INTN.specify.returt.type.in.annotation"); + return true; + } + } PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); @@ -96,9 +110,12 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); + PyParameter parameter = null; + PsiReference reference = null; + String name = null; if (problemElement != null) { - String name = problemElement.getName(); - PsiReference reference = problemElement.getReference(); + name = problemElement.getName(); + reference = problemElement.getReference(); if (problemElement instanceof PyQualifiedExpression) { final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); if (qualifier != null) { @@ -106,70 +123,76 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { name = qualifier.getText(); } } - PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + } + final PsiElement resolvedReference = reference != null? reference.resolve() : null; + if (problemElement instanceof PyParameter) + parameter = (PyParameter)problemElement; + else { + if (resolvedReference instanceof PyParameter) { + parameter = (PyParameter)resolvedReference; + } + } - PyParameter parameter = null; - final PsiElement resolvedReference = reference != null? reference.resolve() : null; - if (problemElement instanceof PyParameter) - parameter = (PyParameter)problemElement; + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + if (parameter != null && name != null) { + final PyFunction function = + elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class, + "def foo(" + name + ": object):\n\tpass"); + final PyNamedParameter namedParameter = function.getParameterList().findParameterByName(name); + assert namedParameter != null; + parameter = (PyParameter)parameter.replace(namedParameter); + parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter); + editor.getCaretModel().moveToOffset(parameter.getTextOffset()); + + final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter); + builder.replaceRange(TextRange.create(parameter.getTextLength()-PyNames.OBJECT.length(), parameter.getTextLength()), PyNames.OBJECT); + Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + TemplateManager.getInstance(project).startTemplate(editor, template); + } + else { //return type + Callable callable = null; + if (resolvedReference instanceof PyTargetExpression) { + final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue(); + if (assignedValue instanceof PyCallExpression) { + callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext()); + } + } else { - if (resolvedReference instanceof PyParameter) { - parameter = (PyParameter)resolvedReference; - } - } - if (parameter != null && name != null) { - final PyFunction function = - elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class, - "def foo(" + name + ": object):\n\tpass"); - final PyNamedParameter namedParameter = function.getParameterList().findParameterByName(name); - assert namedParameter != null; - parameter = (PyParameter)parameter.replace(namedParameter); - parameter = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(parameter); - editor.getCaretModel().moveToOffset(parameter.getTextOffset()); + PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1); + if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER)) + elementAt = file.findElementAt(editor.getCaretModel().getOffset()); - final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter); - builder.replaceRange(TextRange.create(parameter.getTextLength()-PyNames.OBJECT.length(), parameter.getTextLength()), PyNames.OBJECT); + callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + } + if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) { + final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText(); + String functionText = functionSignature + + " -> object:"; + for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) { + functionText = functionText + "\n\t" + st.getText(); + } + final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(callable), PyFunction.class, + functionText); + callable = (PyFunction)callable.replace(function); + callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable); + + final PyExpression value = ((PyFunction)callable).getAnnotation().getValue(); + final int offset = value.getTextOffset(); + + final TemplateBuilder builder = TemplateBuilderFactory.getInstance(). + createTemplateBuilder(value); + builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT); Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - TemplateManager.getInstance(project).startTemplate(editor, template); - } - else { //return type - if (resolvedReference instanceof PyTargetExpression) { - final PyExpression assignedValue = ((PyTargetExpression)resolvedReference).findAssignedValue(); - if (assignedValue instanceof PyCallExpression) { - Callable callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(PyResolveContext.defaultContext()); - if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) { - final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText(); - String functionText = functionSignature + - " -> object:"; - for (PyStatement st : ((PyFunction)callable).getStatementList().getStatements()) { - functionText = functionText + "\n\t" + st.getText(); - } - final PyFunction function = elementGenerator.createFromText(LanguageLevel.forElement(problemElement), PyFunction.class, - functionText); - callable = (PyFunction)callable.replace(function); - callable = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(callable); - - final PyExpression value = ((PyFunction)callable).getAnnotation().getValue(); - final int offset = value.getTextOffset(); - - final TemplateBuilder builder = TemplateBuilderFactory.getInstance(). - createTemplateBuilder(value); - builder.replaceRange(TextRange.create(0, PyNames.OBJECT.length()), PyNames.OBJECT); - Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); - OpenFileDescriptor descriptor = new OpenFileDescriptor( - project, - value.getContainingFile().getVirtualFile(), - offset - ); - Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); - if (targetEditor != null) { - targetEditor.getCaretModel().moveToOffset(offset); - TemplateManager.getInstance(project).startTemplate(targetEditor, template); - } - } - } + OpenFileDescriptor descriptor = new OpenFileDescriptor( + project, + value.getContainingFile().getVirtualFile(), + offset + ); + Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true); + if (targetEditor != null) { + targetEditor.getCaretModel().moveToOffset(offset); + TemplateManager.getInstance(project).startTemplate(targetEditor, template); } - } } } diff --git a/python/testData/intentions/afterTypeInDocstring1.py b/python/testData/intentions/afterTypeInDocstring1.py new file mode 100644 index 000000000000..3422432d15c7 --- /dev/null +++ b/python/testData/intentions/afterTypeInDocstring1.py @@ -0,0 +1,5 @@ +def foo(a, b): + """ + :rtype : object + """ + b = 1 \ No newline at end of file diff --git a/python/testData/intentions/beforeTypeInDocstring1.py b/python/testData/intentions/beforeTypeInDocstring1.py new file mode 100644 index 000000000000..96620cb80b40 --- /dev/null +++ b/python/testData/intentions/beforeTypeInDocstring1.py @@ -0,0 +1,2 @@ +def foo(a, b): + b = 1 \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/PyIntentionTest.java index 3ef18d015a58..2e0a3d5e7fbe 100644 --- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java @@ -237,6 +237,10 @@ public class PyIntentionTest extends PyTestCase { doTest(PyBundle.message("INTN.specify.type")); } + public void testTypeInDocstring1() { + doTest(PyBundle.message("INTN.specify.return.type")); + } + public void testTypeInPy3Annotation() { //PY-7045 doTest(PyBundle.message("INTN.specify.type.in.annotation"), LanguageLevel.PYTHON32); }