Assist in specifying types via docstrings

(Intention for docstring and intention for assertion)
This commit is contained in:
Ekaterina Tuzova
2012-07-20 14:52:58 +04:00
parent 8be4f7c57e
commit 2e7eba4aff
12 changed files with 394 additions and 79 deletions
@@ -127,6 +127,16 @@
<category>Python</category>
</intentionAction>
<intentionAction>
<className>com.jetbrains.python.codeInsight.intentions.SpecifyTypeInDocstringIntention</className>
<category>Python</category>
</intentionAction>
<intentionAction>
<className>com.jetbrains.python.codeInsight.intentions.TypeAssertionIntention</className>
<category>Python</category>
</intentionAction>
<intentionAction>
<className>com.jetbrains.python.codeInsight.intentions.ImportFromToImportIntention</className>
<category>Python</category>
@@ -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.
@@ -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();
}
}
@@ -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;
@@ -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<String, Integer> 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;
}
}
@@ -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;
}
}
@@ -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<String, Integer> 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<String, Integer>(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<String, Integer> 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<String, Integer>(replacementText.toString(), offset);
}
private static class RaiseVisitor extends PyRecursiveElementVisitor {
private boolean myHasRaise = false;
private boolean myHasReturn = false;
@@ -0,0 +1,4 @@
def foo(a, b):
assert isinstance(a, object)
a.
b = 1
@@ -0,0 +1,6 @@
def foo(a, b):
"""
:type a: object
"""
a.
b = 1
@@ -0,0 +1,3 @@
def foo(a, b):
a<caret>.
b = 1
@@ -0,0 +1,3 @@
def foo(a, b):
a<caret>.
b = 1
@@ -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;