From 2e7eba4aff4f98a2469c436aa9eaa3e82de26489 Mon Sep 17 00:00:00 2001 From: Ekaterina Tuzova Date: Fri, 20 Jul 2012 14:52:58 +0400 Subject: [PATCH] Assist in specifying types via docstrings (Intention for docstring and intention for assertion) --- python/src/META-INF/python-plugin-common.xml | 10 ++ .../com/jetbrains/python/PyBundle.properties | 6 + .../python/actions/DocstringQuickFix.java | 88 ++--------- .../intentions/PyDocStubIntention.java | 8 +- .../SpecifyTypeInDocstringIntention.java | 145 ++++++++++++++++++ .../intentions/TypeAssertionIntention.java | 101 ++++++++++++ .../PythonDocumentationProvider.java | 91 ++++++++++- .../testData/intentions/afterTypeAssertion.py | 4 + .../intentions/afterTypeInDocstring.py | 6 + .../intentions/beforeTypeAssertion.py | 3 + .../intentions/beforeTypeInDocstring.py | 3 + .../com/jetbrains/python/PyIntentionTest.java | 8 + 12 files changed, 394 insertions(+), 79 deletions(-) create mode 100644 python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java create mode 100644 python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java create mode 100644 python/testData/intentions/afterTypeAssertion.py create mode 100644 python/testData/intentions/afterTypeInDocstring.py create mode 100644 python/testData/intentions/beforeTypeAssertion.py create mode 100644 python/testData/intentions/beforeTypeInDocstring.py diff --git a/python/src/META-INF/python-plugin-common.xml b/python/src/META-INF/python-plugin-common.xml index 8cf8e8428738..ca04dcd5d281 100644 --- a/python/src/META-INF/python-plugin-common.xml +++ b/python/src/META-INF/python-plugin-common.xml @@ -127,6 +127,16 @@ Python + + com.jetbrains.python.codeInsight.intentions.SpecifyTypeInDocstringIntention + Python + + + + com.jetbrains.python.codeInsight.intentions.TypeAssertionIntention + Python + + com.jetbrains.python.codeInsight.intentions.ImportFromToImportIntention Python diff --git a/python/src/com/jetbrains/python/PyBundle.properties b/python/src/com/jetbrains/python/PyBundle.properties index fd8828138357..9f3bc849bdd6 100644 --- a/python/src/com/jetbrains/python/PyBundle.properties +++ b/python/src/com/jetbrains/python/PyBundle.properties @@ -182,6 +182,12 @@ INTN.transform.into.if.else.statement=Transform conditional expression into if/e # PyConvertTripleQuotedStringIntention INTN.doc.string.stub=Insert documentation string stub +#SpecifyTypeInDocstringIntention +INTN.specify.type=Specify type for reference in docstring + +#TypeAssertionIntention +INTN.insert.assertion=Insert type assertion + # Conflict checker CONFLICT.name.$0.obscured=Name ''{0}'' obscured by local definitions CONFLICT.name.$0.obscured.cannot.convert=Name ''{0}'' obscured. Cannot convert. diff --git a/python/src/com/jetbrains/python/actions/DocstringQuickFix.java b/python/src/com/jetbrains/python/actions/DocstringQuickFix.java index 6dff34a4442f..3054dccb9981 100644 --- a/python/src/com/jetbrains/python/actions/DocstringQuickFix.java +++ b/python/src/com/jetbrains/python/actions/DocstringQuickFix.java @@ -8,9 +8,7 @@ import com.intellij.openapi.editor.EditorFactory; import com.intellij.openapi.project.Project; import com.intellij.openapi.util.text.LineTokenizer; import com.intellij.psi.PsiDocumentManager; -import com.intellij.psi.PsiElement; import com.intellij.psi.PsiFile; -import com.intellij.psi.PsiWhiteSpace; import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.PyBundle; import com.jetbrains.python.documentation.PyDocumentationSettings; @@ -64,7 +62,9 @@ public class DocstringQuickFix implements LocalQuickFix { private static Editor getEditor(Project project, PsiFile file) { Document document = PsiDocumentManager.getInstance(project).getDocument(file); if (document != null) { - Editor[] editors = EditorFactory.getInstance().getEditors(document); + final EditorFactory instance = EditorFactory.getInstance(); + if (instance == null) return null; + Editor[] editors = instance.getEditors(document); if (editors.length > 0) return editors[0]; } @@ -74,14 +74,12 @@ public class DocstringQuickFix implements LocalQuickFix { public void applyFix(@NotNull Project project, @NotNull ProblemDescriptor descriptor) { PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PyDocStringOwner.class); if (docStringOwner == null) return; - PyStringLiteralExpression element = docStringOwner.getDocStringExpression(); - if (element == null && myMissing == null && myUnexpected == null) { + PyStringLiteralExpression docStringExpression = docStringOwner.getDocStringExpression(); + if (docStringExpression == null && myMissing == null && myUnexpected == null) { if (docStringOwner instanceof PyFunction) { - PsiDocumentManager.getInstance(project).getDocument(docStringOwner.getContainingFile()); PythonDocumentationProvider.inserDocStub((PyFunction)docStringOwner, project, getEditor(project, docStringOwner.getContainingFile())); } if (docStringOwner instanceof PyClass) { - PsiDocumentManager.getInstance(project).getDocument(docStringOwner.getContainingFile()); PyFunction init = ((PyClass)docStringOwner).findInitOrNew(false); if (init == null) return; PythonDocumentationProvider.inserDocStub(init, ((PyClass)docStringOwner).getStatementList(), @@ -89,25 +87,26 @@ public class DocstringQuickFix implements LocalQuickFix { } return; } + if (docStringExpression == null) return; PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); - PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(element.getProject()); - if (documentationSettings.isEpydocFormat(element.getContainingFile())) { + PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project); + if (documentationSettings.isEpydocFormat(docStringExpression.getContainingFile())) { myPrefix = "@"; } else { myPrefix = ":"; } - String replacement = element.getText(); + String replacement = docStringExpression.getText(); if (myMissing != null) { - replacement = createMissingReplacement(element); + replacement = createMissingReplacement(docStringOwner); } if (myUnexpected != null) { replacement = createUnexpectedReplacement(replacement); } - if (!replacement.equals(element.getText())) { + if (!replacement.equals(docStringExpression.getText())) { PyExpression str = elementGenerator.createDocstring(replacement).getExpression(); - element.replace(str); + docStringExpression.replace(str); } } @@ -145,66 +144,7 @@ public class DocstringQuickFix implements LocalQuickFix { return newText.toString(); } - private String createMissingReplacement(PsiElement element) { - String text = element.getText(); - String[] lines = LineTokenizer.tokenize(text, true); - StringBuilder replacementText = new StringBuilder(); - int ind = lines.length - 1; - if (lines.length == 1) { - return createSingleLineReplacement(element); - } - for (int i = 0; i != lines.length - 1; ++i) { - String line = lines[i]; - if (line.contains(myPrefix)) { - ind = i; - break; - } - replacementText.append(line); - } - addParam(replacementText, element, false); - for (int i = ind; i != lines.length; ++i) { - String line = lines[i]; - replacementText.append(line); - } - return replacementText.toString(); - } - - private void addParam(StringBuilder replacementText, PsiElement element, boolean addWS) { - PyFunction fun = PsiTreeUtil.getParentOfType(element, PyFunction.class); - PsiWhiteSpace whitespace = PsiTreeUtil.getPrevSiblingOfType(fun.getStatementList(), PsiWhiteSpace.class); - String ws = "\n"; - if (whitespace != null) { - String[] spaces = whitespace.getText().split("\n"); - if (spaces.length > 1) { - ws = ws + whitespace.getText().split("\n")[1]; - } - } - if (replacementText.length() > 0) - replacementText.deleteCharAt(replacementText.length() - 1); - replacementText.append(ws); - - String paramText = myMissingText; - replacementText.append(myPrefix).append("param ").append(paramText).append(": "); - if (addWS) - replacementText.append(ws); - else - replacementText.append("\n"); - } - - private String createSingleLineReplacement(PsiElement element) { - String text = element.getText(); - StringBuilder replacementText = new StringBuilder(); - String closingQuotes = ""; - if (text.endsWith("'''") || text.endsWith("\"\"\"")) { - replacementText.append(text.substring(0, text.length() - 2)); - closingQuotes = text.substring(text.length() - 3); - } - else { - replacementText.append(text.substring(0, text.length())); - closingQuotes = text.substring(text.length() - 1); - } - addParam(replacementText, element, true); - replacementText.append(closingQuotes); - return replacementText.toString(); + private String createMissingReplacement(PyDocStringOwner docstring) { + return PythonDocumentationProvider.addParamToDocstring(docstring, "param", myMissingText, myPrefix).getFirst(); } } diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java index ddea272dc7d7..9f35b489520e 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyDocStubIntention.java @@ -11,6 +11,7 @@ import com.jetbrains.python.PyBundle; import com.jetbrains.python.documentation.PythonDocumentationProvider; import com.jetbrains.python.psi.PyDocStringOwner; import com.jetbrains.python.psi.PyFunction; +import com.jetbrains.python.psi.PyStatementList; import org.jetbrains.annotations.NotNull; /** @@ -33,13 +34,16 @@ public class PyDocStubIntention extends BaseIntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { PyFunction function = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyFunction.class); - if (function != null) { + PyStatementList list = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyStatementList.class, + false, PyFunction.class); + if (function != null && list == null) { final PyDocStringOwner docStringOwner = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), PyDocStringOwner.class); if (docStringOwner != null) { if (docStringOwner.getDocStringExpression() != null) return false; } - if (function.getStatementList() != null && function.getStatementList().getStatements().length != 0) + final PyStatementList statementList = function.getStatementList(); + if (statementList != null && statementList.getStatements().length != 0) return true; } return false; diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java new file mode 100644 index 000000000000..b01b2ad6c93c --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -0,0 +1,145 @@ +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.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.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.documentation.PyDocumentationSettings; +import com.jetbrains.python.documentation.PythonDocumentationProvider; +import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.types.PyReturnTypeReference; +import com.jetbrains.python.psi.types.PyType; +import com.jetbrains.python.psi.types.TypeEvalContext; +import org.jetbrains.annotations.NotNull; + +/** + * User: ktisha + * + * Helps to specify type + */ +public class SpecifyTypeInDocstringIntention implements IntentionAction { + + public SpecifyTypeInDocstringIntention() { + } + + @NotNull + public String getText() { + return PyBundle.message("INTN.specify.type"); + } + + @NotNull + public String getFamilyName() { + return PyBundle.message("INTN.specify.type"); + } + + public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { + PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1), + PyNamedParameter.class); + if (problemElement == null) + problemElement = PsiTreeUtil.getTopmostParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1), + PyQualifiedExpression.class); + if (problemElement == null) return false; + if (problemElement instanceof PyQualifiedExpression) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) { + problemElement = qualifier; + } + } + final PyType type = problemElement.getType(TypeEvalContext.fast()); + if (type == null || type instanceof PyReturnTypeReference) { + PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class); + PsiReference reference = problemElement.getReference(); + if (problemElement instanceof PyQualifiedExpression) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) reference = qualifier.getReference(); + } + if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) + return true; + } + return false; + } + + public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { + PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1), PyNamedParameter.class); + if (problemElement == null) + problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), PyExpression.class); + if (problemElement != null) { + String name = problemElement.getName(); + PsiReference reference = problemElement.getReference(); + if (problemElement instanceof PyQualifiedExpression) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier != null) { + reference = qualifier.getReference(); + name = qualifier.getText(); + } + } + PyDocumentationSettings documentationSettings = PyDocumentationSettings.getInstance(project); + String prefix = ":"; + if (documentationSettings.isEpydocFormat(file)) { + prefix = "@"; + } + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + + PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class); + + 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 { + PyExpressionStatement str = elementGenerator.createDocstring(replacementToOffset.getFirst()); + final PyStatementList list = pyFunction.getStatementList(); + if (list != null && list.getStatements().length != 0) { + list.addBefore(str, list.getStatements()[0]); + } + + pyFunction = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(pyFunction); + docStringExpression = pyFunction.getDocStringExpression(); + startOffset = replacementToOffset.getSecond(); + 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); + } + } + } + + public boolean startInWriteAction() { + return true; + } +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java new file mode 100644 index 000000000000..8806f9a3dbbf --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java @@ -0,0 +1,101 @@ +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.openapi.editor.Editor; +import com.intellij.openapi.project.Project; +import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiElement; +import com.intellij.psi.PsiFile; +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.psi.*; +import com.jetbrains.python.psi.types.PyReturnTypeReference; +import com.jetbrains.python.psi.types.PyType; +import com.jetbrains.python.psi.types.TypeEvalContext; +import org.jetbrains.annotations.NotNull; + +/** + * User: ktisha + * + * Helps to specify type by assertion + */ +public class TypeAssertionIntention implements IntentionAction { + + public TypeAssertionIntention() { + } + + @NotNull + public String getText() { + return PyBundle.message("INTN.insert.assertion"); + } + + @NotNull + public String getFamilyName() { + return PyBundle.message("INTN.insert.assertion"); + } + + public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { + PyExpression problemElement = + PsiTreeUtil.getTopmostParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), PyQualifiedExpression.class); + if (problemElement == null) return false; + if (problemElement instanceof PyQualifiedExpression && ((PyQualifiedExpression)problemElement).getReferencedName() == null) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) { + problemElement = qualifier; + } + } + final PyType type = problemElement.getType(TypeEvalContext.fast()); + return (type == null || type instanceof PyReturnTypeReference); + } + + public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { + PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), + PyExpression.class); + if (problemElement != null) { + PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); + + String name = problemElement.getText(); + if (problemElement instanceof PyQualifiedExpression) { + final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); + if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) { + final String referencedName = ((PyQualifiedExpression)problemElement).getReferencedName(); + if (referencedName == null || PyNames.GETITEM.equals(referencedName)) + name = qualifier.getText(); + } + } + + final String text = "assert isinstance(" + name + ", )"; + PyAssertStatement assertStatement = elementGenerator.createFromText(LanguageLevel.forElement(problemElement), + PyAssertStatement.class, text); + + final PsiElement parentStatement = PsiTreeUtil.getParentOfType(problemElement, PyStatement.class); + if (parentStatement == null) return; + final PsiElement parent = parentStatement.getParent(); + PsiElement element; + if (parentStatement instanceof PyAssignmentStatement && + ((PyAssignmentStatement)parentStatement).getTargets()[0] == problemElement) { + element = parent.addAfter(assertStatement, parentStatement); + } + else { + element = parent.addBefore(assertStatement, parentStatement); + } + + int textOffSet = element.getTextOffset(); + editor.getCaretModel().moveToOffset(textOffSet); + + element = CodeInsightUtilBase.forcePsiPostprocessAndRestoreElement(element); + final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(element); + builder.replaceRange(TextRange.create(text.length()-1, text.length()-1), PyNames.OBJECT); + Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate(); + TemplateManager.getInstance(project).startTemplate(editor, template); + } + } + + public boolean startInWriteAction() { + return true; + } +} \ No newline at end of file diff --git a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java index 4785b7fd8d8a..8ed5fc982957 100644 --- a/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java +++ b/python/src/com/jetbrains/python/documentation/PythonDocumentationProvider.java @@ -11,6 +11,8 @@ import com.intellij.openapi.project.Project; import com.intellij.openapi.projectRoots.Sdk; import com.intellij.openapi.roots.ProjectRootManager; import com.intellij.openapi.ui.Messages; +import com.intellij.openapi.util.Pair; +import com.intellij.openapi.util.text.LineTokenizer; import com.intellij.openapi.util.text.StringUtil; import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; @@ -34,7 +36,10 @@ import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; import java.io.IOException; -import java.util.*; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import static com.jetbrains.python.documentation.DocumentationBuilderKit.*; @@ -473,7 +478,7 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i if (whitespace != null) { String[] spaces = whitespace.getText().split("\n"); if (spaces.length > 1) - ws = ws + spaces[spaces.length-1]; + ws += spaces[spaces.length - 1]; } String docContent = ws + generateDocumentationContentStub(function, ws, true); PyExpressionStatement string = elementGenerator.createDocstring("\"\"\"" + docContent + "\"\"\""); @@ -497,7 +502,7 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i if (whitespace != null) { String[] spaces = whitespace.getText().split("\n"); if (spaces.length > 1) - ws = ws + whitespace.getText().split("\n")[1]; + ws += whitespace.getText().split("\n")[1]; } return generateDocumentationContentStub(element, ws, checkReturn); } @@ -545,6 +550,86 @@ public class PythonDocumentationProvider extends AbstractDocumentationProvider i return builder.toString(); } + public static Pair addParamToDocstring(PyDocStringOwner function, String keyword, String paramName, String prefix) { + final PyStringLiteralExpression docstring = function.getDocStringExpression(); + String text; + if (docstring != null) + text = docstring.getText(); + else + text = "\"\"\"\"\"\""; + String[] lines = LineTokenizer.tokenize(text, true); + StringBuilder replacementText = new StringBuilder(); + int ind = lines.length - 1; + if (lines.length == 1) { + return createSingleLineReplacement(function, keyword, paramName, prefix); + } + for (int i = 0; i != lines.length - 1; ++i) { + String line = lines[i]; + if (line.contains(prefix)) { + ind = i; + break; + } + replacementText.append(line); + } + int offset = addParamOrType(replacementText, function, false, keyword, paramName, prefix); + for (int i = ind; i != lines.length; ++i) { + String line = lines[i]; + replacementText.append(line); + } + return new Pair(replacementText.toString(), offset); + } + + private static int addParamOrType(StringBuilder replacementText, PyDocStringOwner function, boolean addWS, String keyword, + String paramName, String prefix) { + PsiWhiteSpace whitespace = null; + if (function instanceof PyFunction) + whitespace = PsiTreeUtil.getPrevSiblingOfType(((PyFunction)function).getStatementList(), PsiWhiteSpace.class); + String ws = "\n"; + if (whitespace != null) { + String[] spaces = whitespace.getText().split("\n"); + if (spaces.length > 1) { + ws += whitespace.getText().split("\n")[1]; + } + } + if (replacementText.length() > 0) + replacementText.deleteCharAt(replacementText.length() - 1); + replacementText.append(ws); + + replacementText.append(prefix); + replacementText.append(keyword); + replacementText.append(" "); + replacementText.append(paramName); + replacementText.append(": "); + int offset = replacementText.length(); + if (addWS) + replacementText.append(ws); + else + replacementText.append("\n"); + return offset; + } + + private static Pair createSingleLineReplacement(PyDocStringOwner function, String keyword, String paramName, String prefix) { + String text; + final PyStringLiteralExpression docstring = function.getDocStringExpression(); + if (docstring != null) + text = docstring.getText(); + else + text = "\"\"\"\"\"\""; + StringBuilder replacementText = new StringBuilder(); + String closingQuotes = ""; + if (text.endsWith("'''") || text.endsWith("\"\"\"")) { + replacementText.append(text.substring(0, text.length() - 2)); + closingQuotes = text.substring(text.length() - 3); + } + else { + replacementText.append(text.substring(0, text.length())); + closingQuotes = text.substring(text.length() - 1); + } + final int offset = addParamOrType(replacementText, function, true, keyword, paramName, prefix); + replacementText.append(closingQuotes); + return new Pair(replacementText.toString(), offset); + } + private static class RaiseVisitor extends PyRecursiveElementVisitor { private boolean myHasRaise = false; private boolean myHasReturn = false; diff --git a/python/testData/intentions/afterTypeAssertion.py b/python/testData/intentions/afterTypeAssertion.py new file mode 100644 index 000000000000..25419dd10530 --- /dev/null +++ b/python/testData/intentions/afterTypeAssertion.py @@ -0,0 +1,4 @@ +def foo(a, b): + assert isinstance(a, object) + a. + b = 1 \ No newline at end of file diff --git a/python/testData/intentions/afterTypeInDocstring.py b/python/testData/intentions/afterTypeInDocstring.py new file mode 100644 index 000000000000..910a0afdb11f --- /dev/null +++ b/python/testData/intentions/afterTypeInDocstring.py @@ -0,0 +1,6 @@ +def foo(a, b): + """ + :type a: object + """ + a. + b = 1 \ No newline at end of file diff --git a/python/testData/intentions/beforeTypeAssertion.py b/python/testData/intentions/beforeTypeAssertion.py new file mode 100644 index 000000000000..d34e3f2d3c92 --- /dev/null +++ b/python/testData/intentions/beforeTypeAssertion.py @@ -0,0 +1,3 @@ +def foo(a, b): + a. + b = 1 \ No newline at end of file diff --git a/python/testData/intentions/beforeTypeInDocstring.py b/python/testData/intentions/beforeTypeInDocstring.py new file mode 100644 index 000000000000..d34e3f2d3c92 --- /dev/null +++ b/python/testData/intentions/beforeTypeInDocstring.py @@ -0,0 +1,3 @@ +def foo(a, b): + a. + 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 7389a52b9c61..33a0ecb0f03c 100644 --- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java @@ -226,6 +226,14 @@ public class PyIntentionTest extends PyTestCase { doTest("Convert to 'import sys'"); } + public void testTypeInDocstring() { + doTest(PyBundle.message("INTN.specify.type")); + } + + public void testTypeAssertion() { + doTest(PyBundle.message("INTN.insert.assertion")); + } + public void testDocStub() { CodeInsightSettings codeInsightSettings = CodeInsightSettings.getInstance(); codeInsightSettings.JAVADOC_STUB_ON_ENTER = true;