fixed PY-7354 "Add return type annotation" quickfix on function to add :rtype: to docstring

This commit is contained in:
Ekaterina Tuzova
2012-09-03 14:39:21 -04:00
parent b71c670c07
commit cc415992b2
5 changed files with 80 additions and 15 deletions
@@ -20,6 +20,7 @@ import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.documentation.PyDocumentationSettings;
import com.jetbrains.python.documentation.PythonDocumentationProvider;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.resolve.PyResolveContext;
import com.jetbrains.python.psi.types.PyReturnTypeReference;
import com.jetbrains.python.psi.types.PyType;
import com.jetbrains.python.psi.types.TypeEvalContext;
@@ -50,12 +51,25 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
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;
PyCallExpression callExpression = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class);
if (callExpression != null ) {
PyAssignmentStatement assignmentStatement = PsiTreeUtil.getParentOfType(elementAt, PyAssignmentStatement.class);
if (assignmentStatement != null) {
PyType type = assignmentStatement.getAssignedValue().getType(TypeEvalContext.slow());
if (type == null || type instanceof PyReturnTypeReference) {
myText = PyBundle.message("INTN.specify.return.type");
return true;
}
}
}
else {
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;
}
}
}
@@ -104,7 +118,20 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
String type = "type";
String name = "";
PyCallExpression callExpression = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class);
PyFunction pyFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
if (callExpression != null ) {
PyAssignmentStatement assignmentStatement = PsiTreeUtil.getParentOfType(elementAt, PyAssignmentStatement.class);
if (assignmentStatement != null) {
PyType pyType = assignmentStatement.getAssignedValue().getType(TypeEvalContext.slow());
if (pyType == null || pyType instanceof PyReturnTypeReference) {
pyFunction = (PyFunction)callExpression.resolveCalleeFunction(PyResolveContext.defaultContext());
problemElement = null;
type = "rtype";
}
}
}
if (pyFunction != null) {
final ASTNode nameNode = pyFunction.getNameNode();
if (nameNode != null && nameNode.getPsi() == elementAt) {
@@ -119,7 +146,7 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
}
PsiReference reference = null;
PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class);
PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);
if (problemElement != null) {
name = problemElement.getName();
@@ -137,7 +164,7 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction {
final ASTNode nameNode = pyFunction.getNameNode();
if ((pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) ||
elementAt == nameNode.getPsi()) {
elementAt == nameNode.getPsi() || callExpression != null) {
PyStringLiteralExpression docStringExpression = pyFunction.getDocStringExpression();
int startOffset;
int endOffset;
@@ -50,12 +50,25 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
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;
PyCallExpression callExpression = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class);
if (callExpression != null ) {
PyAssignmentStatement assignmentStatement = PsiTreeUtil.getParentOfType(elementAt, PyAssignmentStatement.class);
if (assignmentStatement != null) {
PyType type = assignmentStatement.getAssignedValue().getType(TypeEvalContext.slow());
if (type == null || type instanceof PyReturnTypeReference) {
myText = PyBundle.message("INTN.specify.returt.type.in.annotation");
return true;
}
}
}
else {
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;
}
}
}
@@ -159,10 +172,16 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction {
}
else {
PsiElement elementAt = file.findElementAt(editor.getCaretModel().getOffset() - 1);
PyCallExpression callExpression = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class);
if (elementAt != null && !(elementAt.getNode().getElementType() == PyTokenTypes.IDENTIFIER))
elementAt = file.findElementAt(editor.getCaretModel().getOffset());
callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
if (callExpression != null) {
callable = callExpression.resolveCalleeFunction(PyResolveContext.defaultContext());
}
else
callable = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class);
}
if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() == null) {
final String functionSignature = "def " + callable.getName() + callable.getParameterList().getText();
@@ -0,0 +1,9 @@
def func1(x):
"""
:rtype : object
"""
return x
def func2(x):
y = func1(x.keys())
return y.startswith('foo')
@@ -0,0 +1,6 @@
def func1(x):
return x
def func2(x):
y = fu<caret>nc1(x.keys())
return y.startswith('foo')
@@ -241,6 +241,10 @@ public class PyIntentionTest extends PyTestCase {
doTest(PyBundle.message("INTN.specify.return.type"));
}
public void testTypeInDocstring2() {
doTest(PyBundle.message("INTN.specify.return.type"));
}
public void testTypeInPy3Annotation() { //PY-7045
doTest(PyBundle.message("INTN.specify.type.in.annotation"), LanguageLevel.PYTHON32);
}