diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java index e94104501c6b..21ad7860d16b 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java @@ -18,7 +18,6 @@ package com.jetbrains.python.codeInsight.intentions; import com.google.common.collect.Lists; import com.intellij.codeInsight.CodeInsightUtilCore; import com.intellij.codeInsight.FileModificationService; -import com.intellij.codeInsight.TargetElementUtil; import com.intellij.codeInsight.template.*; import com.intellij.openapi.application.WriteAction; import com.intellij.openapi.editor.Document; @@ -55,11 +54,7 @@ public class PyAnnotateTypesIntention extends PyBaseIntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false; - final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); - if (elementAt == null) return false; - - final PyFunction function = findSuitableFunction(elementAt); + final PyFunction function = findSuitableFunction(editor, file); if (function != null) { setText(PyBundle.message("INTN.add.type.hints.for.function", function.getName())); return true; @@ -67,9 +62,12 @@ public class PyAnnotateTypesIntention extends PyBaseIntentionAction { return false; } - @Nullable - public PyFunction findSuitableFunction(@NotNull PsiElement elementAt) { - return TypeIntention.findSuitableFunction(elementAt, input -> true); + @Override + public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { + final PyFunction function = findSuitableFunction(editor, file); + if (function != null) { + annotateTypes(editor, function); + } } @Override @@ -77,16 +75,9 @@ public class PyAnnotateTypesIntention extends PyBaseIntentionAction { return false; } - @Override - public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); - if (elementAt != null) { - final PyFunction function = findSuitableFunction(elementAt); - if (function != null) { - annotateTypes(editor, function); - } - } + @Nullable + private static PyFunction findSuitableFunction(@NotNull Editor editor, @NotNull PsiFile file) { + return TypeIntention.findOnlySuitableFunction(editor, file, input -> true); } public static void annotateTypes(Editor editor, PyFunction function) { diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java index 5658e7276b87..bee4db9f2796 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -16,14 +16,11 @@ package com.jetbrains.python.codeInsight.intentions; import com.intellij.codeInsight.FileModificationService; -import com.intellij.codeInsight.TargetElementUtil; import com.intellij.openapi.application.WriteAction; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.project.Project; import com.intellij.openapi.util.text.StringUtil; -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.intellij.util.ObjectUtils; @@ -33,7 +30,9 @@ import com.jetbrains.python.debugger.PySignature; import com.jetbrains.python.debugger.PySignatureCacheManager; import com.jetbrains.python.documentation.docstrings.DocStringUtil; import com.jetbrains.python.documentation.docstrings.PyDocstringGenerator; -import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.PyFunction; +import com.jetbrains.python.psi.PyNamedParameter; +import com.jetbrains.python.psi.StructuredDocString; import com.jetbrains.python.toolbox.Substring; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -58,25 +57,18 @@ public class SpecifyTypeInDocstringIntention extends TypeIntention { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); - final PyExpression problemElement = getProblemElement(elementAt); - final PsiReference reference = problemElement == null ? null : problemElement.getReference(); - - final PsiElement resolved = reference != null ? reference.resolve() : null; - final PyNamedParameter parameter = getParameter(problemElement, resolved); - + final PyNamedParameter parameter = findOnlySuitableParameter(editor, file); if (parameter != null) { final PyFunction parentFunction = PsiTreeUtil.getParentOfType(parameter, PyFunction.class); if (parentFunction != null) { generateDocstring(parameter, parentFunction); } + return; } - else if (elementAt != null) { - final PyFunction function = findSuitableFunction(elementAt); - if (function != null) { - generateDocstring(null, function); - } + + final PyFunction function = findOnlySuitableFunction(editor, file); + if (function != null) { + generateDocstring(null, function); } } @@ -120,7 +112,7 @@ public class SpecifyTypeInDocstringIntention extends TypeIntention { } @Override - protected boolean isParamTypeDefined(@NotNull PyParameter parameter) { + protected boolean isParamTypeDefined(@NotNull PyNamedParameter parameter) { final PyFunction pyFunction = PsiTreeUtil.getParentOfType(parameter, PyFunction.class); if (pyFunction != null) { final StructuredDocString structuredDocString = pyFunction.getStructuredDocString(); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index 69fc69bbfa9e..bfa497c9c73d 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -17,7 +17,6 @@ package com.jetbrains.python.codeInsight.intentions; import com.intellij.codeInsight.CodeInsightUtilCore; import com.intellij.codeInsight.FileModificationService; -import com.intellij.codeInsight.TargetElementUtil; import com.intellij.codeInsight.template.*; import com.intellij.openapi.application.WriteAction; import com.intellij.openapi.editor.Editor; @@ -28,7 +27,6 @@ import com.intellij.openapi.util.TextRange; import com.intellij.openapi.util.text.StringUtil; 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.intellij.util.ObjectUtils; @@ -69,22 +67,15 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); - final PyExpression problemElement = getProblemElement(elementAt); - final PsiReference reference = problemElement == null ? null : problemElement.getReference(); - - final PsiElement resolved = reference != null ? reference.resolve() : null; - final PyNamedParameter parameter = getParameter(problemElement, resolved); - + final PyNamedParameter parameter = findOnlySuitableParameter(editor, file); if (parameter != null) { annotateParameter(project, editor, parameter); + return; } - else if (elementAt != null) { - final PyFunction function = findSuitableFunction(elementAt); - if (function != null) { - annotateReturnType(project, function); - } + + final PyFunction function = findOnlySuitableFunction(editor, file); + if (function != null) { + annotateReturnType(project, function); } } @@ -215,16 +206,8 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { } @Override - protected boolean isParamTypeDefined(PyParameter parameter) { - return isDefinedInAnnotation(parameter); - } - - private static boolean isDefinedInAnnotation(PyParameter parameter) { - if (LanguageLevel.forElement(parameter).isPython2()) { - return false; - } - if (parameter instanceof PyNamedParameter && (((PyNamedParameter)parameter).getAnnotation() != null)) return true; - return false; + protected boolean isParamTypeDefined(@NotNull PyNamedParameter parameter) { + return parameter.getAnnotation() != null; } @Override diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java index 17d48e4d7ef8..09fc13dcd08d 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java @@ -20,9 +20,8 @@ import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.project.Project; import com.intellij.openapi.roots.ProjectFileIndex; -import com.intellij.openapi.roots.ProjectRootManager; -import com.intellij.openapi.vfs.VirtualFile; -import com.intellij.psi.*; +import com.intellij.psi.PsiElement; +import com.intellij.psi.PsiFile; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.containers.ContainerUtil; import com.jetbrains.python.documentation.doctest.PyDocstringFile; @@ -33,11 +32,10 @@ import one.util.streamex.StreamEx; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import java.util.Collections; import java.util.List; import java.util.function.Predicate; -import static com.jetbrains.python.psi.PyUtil.as; - /** * User: ktisha * @@ -49,93 +47,83 @@ public abstract class TypeIntention extends PyBaseIntentionAction { if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false; updateText(false); - final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, offset); - if (elementAt == null) return false; - if (isAvailableForParameter(project, elementAt)) { + if (findOnlySuitableParameter(editor, file) != null) { return true; } - if (isAvailableForReturn(elementAt)) { + if (findOnlySuitableFunction(editor, file) != null) { 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) { - return false; - } - final PsiReference reference = problemElement.getReference(); - if (reference instanceof PsiPolyVariantReference) { - final ResolveResult[] results = ((PsiPolyVariantReference)reference).multiResolve(false); - if (results.length != 1) return false; - } - final VirtualFile virtualFile = problemElement.getContainingFile().getVirtualFile(); - if (virtualFile != null) { - if (ProjectRootManager.getInstance(project).getFileIndex().isInLibraryClasses(virtualFile)) { - return false; - } - } - final PsiElement resolved = reference != null ? reference.resolve() : null; - final PyParameter parameter = getParameter(problemElement, resolved); - - return parameter != null && !isParamTypeDefined(parameter); + @Nullable + protected final PyFunction findOnlySuitableFunction(@NotNull Editor editor, @NotNull PsiFile file) { + return findOnlySuitableFunction(editor, file, input -> !isReturnTypeDefined(input)); } @Nullable - public static PyExpression getProblemElement(@Nullable PsiElement elementAt) { - return PsiTreeUtil.getParentOfType(elementAt, PyNamedParameter.class, PyReferenceExpression.class); + public static PyFunction findOnlySuitableFunction(@NotNull Editor editor, @NotNull PsiFile file, Predicate condition) { + final PsiElement elementAt = getElementUnderCaret(editor, file); + return elementAt != null ? ContainerUtil.getOnlyItem(findSuitableFunctions(elementAt, condition)) : null; + } + + @Nullable + protected final PyNamedParameter findOnlySuitableParameter(@NotNull Editor editor, @NotNull PsiFile file) { + return ContainerUtil.getOnlyItem(findSuitableParameters(editor, file)); + } + + @NotNull + private List findSuitableParameters(@NotNull Editor editor, @NotNull PsiFile file) { + final PsiElement elementAt = getElementUnderCaret(editor, file); + final StreamEx parameters; + final PyNamedParameter immediateParam = PsiTreeUtil.getParentOfType(elementAt, PyNamedParameter.class); + if (immediateParam != null) { + parameters = StreamEx.of(immediateParam); + } + else { + final PyReferenceExpression referenceExpr = PsiTreeUtil.getParentOfType(elementAt, PyReferenceExpression.class); + if (referenceExpr != null) { + parameters = StreamEx.of(PyUtil.multiResolveTopPriority(referenceExpr, getResolveContext(elementAt))) + .select(PyNamedParameter.class); + } + else { + parameters = StreamEx.empty(); + } + } + final ProjectFileIndex index = ProjectFileIndex.getInstance(file.getProject()); + return parameters + .filter(param -> !param.isSelf()) + .filter(param -> PsiTreeUtil.getParentOfType(param, PyLambdaExpression.class) == null) + .filter(param -> !index.isInLibraryClasses(param.getContainingFile().getVirtualFile())) + .filter(param -> !isParamTypeDefined(param)) + .toList(); + } + + @Nullable + private static PsiElement getElementUnderCaret(@NotNull Editor editor, @NotNull PsiFile file) { + final int offset = TargetElementUtil.adjustOffset(file, editor.getDocument(), editor.getCaretModel().getOffset()); + return PyUtil.findNonWhitespaceAtOffset(file, offset); } protected abstract void updateText(boolean isReturn); - protected boolean isParamTypeDefined(PyParameter parameter) { - return false; - } + protected abstract boolean isParamTypeDefined(@NotNull PyNamedParameter parameter); - @Nullable - protected static PyNamedParameter getParameter(PyExpression problemElement, PsiElement resolved) { - PyNamedParameter parameter = as(problemElement, PyNamedParameter.class); - if (resolved instanceof PyNamedParameter) { - parameter = (PyNamedParameter)resolved; - } - return parameter == null || parameter.isSelf() ? null : parameter; - } - - private boolean isAvailableForReturn(@NotNull final PsiElement elementAt) { - return findSuitableFunction(elementAt) != null; - } - - @Nullable - protected PyFunction findSuitableFunction(@NotNull PsiElement elementAt) { - return findSuitableFunction(elementAt, input -> !isReturnTypeDefined(input)); - } - - @Nullable - public static PyFunction findSuitableFunction(@NotNull PsiElement elementAt, @NotNull Predicate extraCondition) { - return ContainerUtil.getOnlyItem(findSuitableFunctions(elementAt, extraCondition)); - } + protected abstract boolean isReturnTypeDefined(@NotNull PyFunction function); @NotNull private static List findSuitableFunctions(@NotNull PsiElement elementAt, @NotNull Predicate extraCondition) { final StreamEx definitions; - final PyFunction underCaret = findFunctionDefinitionUnderCaret(elementAt); - if (underCaret != null) { - definitions = StreamEx.of(underCaret); + final PyFunction immediateDefinition = findFunctionDefinitionUnderCaret(elementAt); + if (immediateDefinition != null) { + definitions = StreamEx.of(immediateDefinition); } else { - final PyCallExpression callExpression = getCallExpression(elementAt); - if (callExpression == null) { - definitions = StreamEx.empty(); - } - else { - definitions = StreamEx.of(callExpression.multiResolveCallee(getResolveContext(elementAt))) - .map(result -> result.getElement()) - .select(PyFunction.class); - } + definitions = StreamEx.of(getCallExpressions(elementAt)) + .flatMap(call -> StreamEx.of(call.multiResolveCallee(getResolveContext(elementAt)))) + .map(result -> result.getElement()) + .select(PyFunction.class); } final ProjectFileIndex index = ProjectFileIndex.getInstance(elementAt.getProject()); return definitions.filter(elem -> !index.isInLibraryClasses(elem.getContainingFile().getVirtualFile())) @@ -158,38 +146,38 @@ public abstract class TypeIntention extends PyBaseIntentionAction { return null; } - protected boolean isReturnTypeDefined(@NotNull PyFunction function) { - return false; - } - - @Nullable - static PyCallExpression getCallExpression(PsiElement elementAt) { - final PyExpression problemElement = getProblemElement(elementAt); - if (problemElement != null) { - PsiReference reference = problemElement.getReference(); - final PsiElement resolved = reference != null? reference.resolve() : null; - if (resolved instanceof PyTargetExpression) { - final PyResolveContext context = getResolveContext(elementAt); - if (context.getTypeEvalContext().maySwitchToAST(resolved)) { - final PyExpression assignedValue = ((PyTargetExpression)resolved).findAssignedValue(); - if (assignedValue instanceof PyCallExpression) { - return (PyCallExpression)assignedValue; - } - } + @NotNull + private static List getCallExpressions(@NotNull PsiElement elementAt) { + final PyResolveContext context = getResolveContext(elementAt); + final PyReferenceExpression referenceExpr = PsiTreeUtil.getParentOfType(elementAt, PyReferenceExpression.class); + if (referenceExpr != null) { + final List calls = StreamEx.of(PyUtil.multiResolveTopPriority(referenceExpr, context)) + .select(PyTargetExpression.class) + .filter(target -> context.getTypeEvalContext().maySwitchToAST(target)) + .map(target -> target.findAssignedValue()) + .select(PyCallExpression.class) + .toList(); + if (!calls.isEmpty()) { + return calls; } } - PyAssignmentStatement assignmentStatement = PsiTreeUtil.getParentOfType(elementAt, PyAssignmentStatement.class); - if (assignmentStatement != null) { - final PyExpression assignedValue = assignmentStatement.getAssignedValue(); + final PyAssignmentStatement assignment = PsiTreeUtil.getParentOfType(elementAt, PyAssignmentStatement.class); + if (assignment != null) { + final PyExpression assignedValue = assignment.getAssignedValue(); if (assignedValue instanceof PyCallExpression) { - return (PyCallExpression)assignedValue; + return Collections.singletonList((PyCallExpression)assignedValue); } } - return PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class, false); + final PyCallExpression immediateCall = PsiTreeUtil.getParentOfType(elementAt, PyCallExpression.class, false); + if (immediateCall != null) { + return Collections.singletonList(immediateCall); + } + return Collections.emptyList(); } - protected static PyResolveContext getResolveContext(@NotNull PsiElement origin) { - return PyResolveContext.noImplicits().withTypeEvalContext(TypeEvalContext.codeAnalysis(origin.getProject(), origin.getContainingFile())); + private static PyResolveContext getResolveContext(@NotNull PsiElement origin) { + final TypeEvalContext typeEvalContext = TypeEvalContext.codeAnalysis(origin.getProject(), origin.getContainingFile()); + return PyResolveContext.noImplicits().withTypeEvalContext(typeEvalContext); } } \ No newline at end of file diff --git a/python/testData/intentions/paramTypeInAnnotationNotSuggestedForLambda.py b/python/testData/intentions/paramTypeInAnnotationNotSuggestedForLambda.py new file mode 100644 index 000000000000..ade2f3020a7a --- /dev/null +++ b/python/testData/intentions/paramTypeInAnnotationNotSuggestedForLambda.py @@ -0,0 +1 @@ +lambda param: 42 \ No newline at end of file diff --git a/python/testData/intentions/paramTypeInAnnotationNotSuggestedForSelf.py b/python/testData/intentions/paramTypeInAnnotationNotSuggestedForSelf.py new file mode 100644 index 000000000000..2caeb2ec6bee --- /dev/null +++ b/python/testData/intentions/paramTypeInAnnotationNotSuggestedForSelf.py @@ -0,0 +1,3 @@ +class MyClass(object): + def method(self): + pass diff --git a/python/testData/intentions/paramTypeInDocstringNotSuggestedForLambda.py b/python/testData/intentions/paramTypeInDocstringNotSuggestedForLambda.py new file mode 100644 index 000000000000..ade2f3020a7a --- /dev/null +++ b/python/testData/intentions/paramTypeInDocstringNotSuggestedForLambda.py @@ -0,0 +1 @@ +lambda param: 42 \ No newline at end of file diff --git a/python/testData/intentions/paramTypeInDocstringNotSuggestedForSelf.py b/python/testData/intentions/paramTypeInDocstringNotSuggestedForSelf.py new file mode 100644 index 000000000000..2caeb2ec6bee --- /dev/null +++ b/python/testData/intentions/paramTypeInDocstringNotSuggestedForSelf.py @@ -0,0 +1,3 @@ +class MyClass(object): + def method(self): + pass diff --git a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java index 003d37722eef..4f6359b66dd4 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyIntentionTest.java @@ -352,6 +352,22 @@ public class PyIntentionTest extends PyTestCase { doDocParamTypeTest(DocStringFormat.REST); } + public void testParamTypeInDocstringNotSuggestedForSelf() { + doNegativeTest(PyBundle.message("INTN.specify.type")); + } + + public void testParamTypeInAnnotationNotSuggestedForSelf() { + runWithLanguageLevel(LanguageLevel.PYTHON36, () -> doNegativeTest(PyBundle.message("INTN.specify.type.in.annotation"))); + } + + public void testParamTypeInDocstringNotSuggestedForLambda() { + doNegativeTest(PyBundle.message("INTN.specify.type")); + } + + public void testParamTypeInAnnotationNotSuggestedForLambda() { + runWithLanguageLevel(LanguageLevel.PYTHON36, () -> doNegativeTest(PyBundle.message("INTN.specify.type.in.annotation"))); + } + // PY-16456 public void testReturnTypeInDocStringDifferentIndentationSize() { doDocReturnTypeTest(DocStringFormat.REST);