diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java index 9304f926d9a3..f7b4909f79c6 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -80,28 +80,26 @@ public class SpecifyTypeInDocstringIntention extends TypeIntention { } @Override - protected boolean isTypeDefined(PyExpression problemElement) { - return isDefinedInDocstring(problemElement); - } - - private boolean isDefinedInDocstring(PyExpression problemElement) { - PsiReference reference = problemElement.getReference(); - PyFunction pyFunction = PsiTreeUtil.getParentOfType(problemElement, PyFunction.class); - if (pyFunction != null && (problemElement instanceof PyParameter || reference != null && reference.resolve() instanceof PyParameter)) { + protected boolean isParamTypeDefined(PyParameter parameter) { + PyFunction pyFunction = PsiTreeUtil.getParentOfType(parameter, PyFunction.class); + if (pyFunction != null && parameter != null) { final String docstring = pyFunction.getDocStringValue(); if (docstring != null) { - String name = problemElement.getName(); - if (problemElement instanceof PyQualifiedExpression) { - final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); - if (qualifier != null) { - name = qualifier.getText(); - } - } StructuredDocString structuredDocString = StructuredDocString.parse(docstring); - return structuredDocString != null && structuredDocString.getParamType(name) != null; + return structuredDocString != null && structuredDocString.getParamType(StringUtil.notNullize(parameter.getName())) != null; } return false; } return false; } + + @Override + protected boolean isReturnTypeDefined(@NotNull PyFunction function) { + final String docstring = function.getDocStringValue(); + if (docstring != null) { + StructuredDocString structuredDocString = StructuredDocString.parse(docstring); + return structuredDocString != null && structuredDocString.getReturnType( ) != null; + } + return false; + } } \ No newline at end of file diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index 437f67a721b4..00976042431b 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -124,37 +124,22 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { } @Override - protected boolean isTypeDefined(PyExpression problemElement) { - - return isDefinedInAnnotation(problemElement); + protected boolean isParamTypeDefined(PyParameter parameter) { + return isDefinedInAnnotation(parameter); } - private boolean isDefinedInAnnotation(PyExpression problemElement) { - if (LanguageLevel.forElement(problemElement).isOlderThan(LanguageLevel.PYTHON30)) { + private boolean isDefinedInAnnotation(PyParameter parameter) { + if (LanguageLevel.forElement(parameter).isOlderThan(LanguageLevel.PYTHON30)) { return false; } - PsiReference reference = problemElement.getReference(); - final PsiElement resolved = reference != null? reference.resolve() : null; - PyParameter parameter = getParameter(problemElement, resolved); - if (parameter instanceof PyNamedParameter && (((PyNamedParameter)parameter).getAnnotation() != null)) return true; - - if (resolved instanceof PyTargetExpression) { // return type - final PyExpression assignedValue = ((PyTargetExpression)resolved).findAssignedValue(); - if (assignedValue instanceof PyCallExpression) { - final PyExpression callee = ((PyCallExpression)assignedValue).getCallee(); - if (callee != null) { - final PsiReference psiReference = callee.getReference(); - if (psiReference != null && psiReference.resolve() == null) return false; - } - final Callable callable = ((PyCallExpression)assignedValue).resolveCalleeFunction(getResolveContext(problemElement)); - - if (callable instanceof PyFunction && ((PyFunction)callable).getAnnotation() != null) return true; - } - } return false; } + @Override + protected boolean isReturnTypeDefined(@NotNull PyFunction function) { + return function.getAnnotation() != null; + } @Override protected void updateText(boolean isReturn) { diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java index 7ca666e00678..5777e44ee12e 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java @@ -1,7 +1,6 @@ package com.jetbrains.python.codeInsight.intentions; import com.intellij.codeInsight.intention.IntentionAction; -import com.intellij.injected.editor.EditorWindow; import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.project.Project; @@ -10,10 +9,9 @@ import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.psi.util.PsiTreeUtil; import com.jetbrains.python.PyNames; +import com.jetbrains.python.documentation.doctest.PyDocstringFile; 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; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -26,16 +24,22 @@ import org.jetbrains.annotations.Nullable; public abstract class TypeIntention implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { - if (editor instanceof EditorWindow) return false; + if (file instanceof PyDocstringFile) return false; updateText(false); - PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); if (elementAt == null) return false; + if (isAvailableForParameter(project, elementAt)) { + return true; + } if (isAvailableForReturn(elementAt)) { updateText(true); return true; } + return false; + } + private boolean isAvailableForParameter(Project project, PsiElement elementAt) { final PyExpression problemElement = getProblemElement(elementAt); if (problemElement == null) return false; if (PsiTreeUtil.getParentOfType(problemElement, PyLambdaExpression.class) != null) { @@ -52,7 +56,10 @@ public abstract class TypeIntention implements IntentionAction { return false; } } - return !isTypeDefined(problemElement); + final PsiElement resolved = reference != null ? reference.resolve() : null; + final PyParameter parameter = getParameter(problemElement, resolved); + + return parameter != null && !isParamTypeDefined(parameter); } @Nullable @@ -70,7 +77,7 @@ public abstract class TypeIntention implements IntentionAction { protected abstract void updateText(boolean isReturn); - protected boolean isTypeDefined(PyExpression problemElement) { + protected boolean isParamTypeDefined(PyParameter parameter) { return false; } @@ -83,44 +90,39 @@ public abstract class TypeIntention implements IntentionAction { } private boolean isAvailableForReturn(PsiElement elementAt) { - PyCallExpression callExpression = getCallExpression(elementAt); - - if (callExpression != null) { - final PyExpression callee = callExpression.getCallee(); - if (callee instanceof PyQualifiedExpression) { - final TypeEvalContext context = TypeEvalContext.fastStubOnly(callee.getContainingFile()); - final Callable pyFunction = callExpression.resolveCalleeFunction(PyResolveContext.defaultContext().withTypeEvalContext(context)); - if (pyFunction instanceof PyFunction) { - PyType type = pyFunction.getReturnType(context, (PyQualifiedExpression)callee); - if (type == null || type instanceof PyReturnTypeReference) { - final PsiReference reference = callee.getReference(); - if (reference instanceof PsiPolyVariantReference) { - final ResolveResult[] results = ((PsiPolyVariantReference)reference).multiResolve(false); - if (results.length == 1) { - final PsiElement result = results[0].getElement(); - if (result == null) return false; - final PsiFile psiFile = result.getContainingFile(); - if (psiFile == null) return false; - final VirtualFile virtualFile = psiFile.getVirtualFile(); - if (virtualFile != null) { - if (ProjectRootManager.getInstance(psiFile.getProject()).getFileIndex().isInLibraryClasses(virtualFile)) { - return false; - } - } - return true; - } - } - } - } - } - } - PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); + final PyFunction parentFunction = PsiTreeUtil.getParentOfType(elementAt, PyFunction.class); if (parentFunction != null) { final ASTNode nameNode = parentFunction.getNameNode(); if (nameNode != null && nameNode.getPsi() == elementAt) { - return true; + return !isReturnTypeDefined(parentFunction); } } + + final PyCallExpression callExpression = getCallExpression(elementAt); + if (callExpression == null) return false; + final PyExpression callee = callExpression.getCallee(); + if (callee == null) return false; + final PsiReference reference = callee.getReference(); + if (reference instanceof PsiPolyVariantReference) { + final ResolveResult[] results = ((PsiPolyVariantReference)reference).multiResolve(false); + if (results.length == 1) { + final PsiElement result = results[0].getElement(); + if (!(result instanceof PyFunction)) return false; + final PsiFile psiFile = result.getContainingFile(); + if (psiFile == null) return false; + final VirtualFile virtualFile = psiFile.getVirtualFile(); + if (virtualFile != null) { + if (ProjectRootManager.getInstance(psiFile.getProject()).getFileIndex().isInLibraryClasses(virtualFile)) { + return false; + } + } + return !isReturnTypeDefined((PyFunction)result); + } + } + return false; + } + + protected boolean isReturnTypeDefined(@NotNull PyFunction function) { return false; } diff --git a/python/testData/intentions/afterReturnTypeInDocstring.py b/python/testData/intentions/afterReturnTypeInDocstring.py new file mode 100644 index 000000000000..d3d1bd206a42 --- /dev/null +++ b/python/testData/intentions/afterReturnTypeInDocstring.py @@ -0,0 +1,12 @@ +def f(x): + """ + + :rtype : object + """ + pass + + +def g(x): + y = x + f(x) # (1) + f(y) # (2) \ No newline at end of file diff --git a/python/testData/intentions/beforeReturnTypeInDocstring.py b/python/testData/intentions/beforeReturnTypeInDocstring.py new file mode 100644 index 000000000000..00dcd8c39b4e --- /dev/null +++ b/python/testData/intentions/beforeReturnTypeInDocstring.py @@ -0,0 +1,8 @@ +def f(x): + pass + + +def g(x): + y = x + f(x) # (1) + f(y) # (2) \ 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 938f988a0559..e9a69cf2cd3e 100644 --- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java @@ -297,12 +297,8 @@ public class PyIntentionTest extends PyTestCase { doDocReferenceTest(); } - private void doDocReferenceTest() { - doTest(PyBundle.message("INTN.specify.type")); - } - - private void doDocReturnTypeTest() { - doTest(PyBundle.message("INTN.specify.return.type")); + public void testReturnTypeInDocstring() { + doDocReturnTypeTest(); } public void testTypeInDocstring1() { @@ -420,4 +416,13 @@ public class PyIntentionTest extends PyTestCase { documentationSettings.setFormat(DocStringFormat.PLAIN); } } + + private void doDocReferenceTest() { + doTest(PyBundle.message("INTN.specify.type")); + } + + private void doDocReturnTypeTest() { + doTest(PyBundle.message("INTN.specify.return.type")); + } + }