diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java index 757a01dffea9..bfd60427b2c5 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -43,11 +43,8 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction { } 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); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); + if (problemElement == null) return false; if (problemElement instanceof PyQualifiedExpression) { final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); @@ -85,9 +82,7 @@ public class SpecifyTypeInDocstringIntention implements IntentionAction { } 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); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); if (problemElement != null) { String name = problemElement.getName(); PsiReference reference = problemElement.getReference(); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index e3ab71f0b1c8..e1e098d08b06 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -45,11 +45,8 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { if (!LanguageLevel.forElement(file).isPy3K()) return false; - 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); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); + if (problemElement == null) return false; if (problemElement instanceof PyQualifiedExpression) { final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); @@ -98,9 +95,7 @@ public class SpecifyTypeInPy3AnnotationsIntention implements IntentionAction { } 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); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyNamedParameter.class, PyQualifiedExpression.class); if (problemElement != null) { String name = problemElement.getName(); PsiReference reference = problemElement.getReference(); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java index 18ae05f5dfb7..2b40bfa7636c 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeAssertionIntention.java @@ -41,8 +41,7 @@ public class TypeAssertionIntention implements IntentionAction { } public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { - PyExpression problemElement = - PsiTreeUtil.getTopmostParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), PyQualifiedExpression.class); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyQualifiedExpression.class); if (problemElement == null) return false; if (problemElement instanceof PyQualifiedExpression) { final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); @@ -63,8 +62,7 @@ public class TypeAssertionIntention implements IntentionAction { } public void invoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - PyExpression problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()-1), - PyExpression.class); + PyExpression problemElement = PyUtil.findProblemElement(editor, file, PyQualifiedExpression.class); if (problemElement != null) { PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project); diff --git a/python/src/com/jetbrains/python/psi/PyUtil.java b/python/src/com/jetbrains/python/psi/PyUtil.java index d9875cf384ad..1a00f6917974 100644 --- a/python/src/com/jetbrains/python/psi/PyUtil.java +++ b/python/src/com/jetbrains/python/psi/PyUtil.java @@ -8,6 +8,7 @@ import com.intellij.codeInsight.lookup.LookupElementBuilder; import com.intellij.ide.fileTemplates.FileTemplate; import com.intellij.ide.fileTemplates.FileTemplateManager; import com.intellij.lang.ASTNode; +import com.intellij.openapi.editor.Editor; import com.intellij.openapi.extensions.Extensions; import com.intellij.openapi.module.Module; import com.intellij.openapi.module.ModuleUtil; @@ -1039,9 +1040,6 @@ public class PyUtil { else if (expression instanceof PySequenceExpression) { valuesLength = ((PySequenceExpression)expression).getElements().length; } - else if (expression instanceof PyDictLiteralExpression) { - valuesLength = ((PyDictLiteralExpression)expression).getElements().length; - } else if (expression instanceof PyStringLiteralExpression) { valuesLength = ((PyStringLiteralExpression)expression).getStringValue().length(); } @@ -1105,5 +1103,18 @@ public class PyUtil { } return null; } + + @Nullable + public static T findProblemElement(Editor editor, PsiFile file, @NotNull final Class... classes) { + for (Class claz : classes) { + PsiElement problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset() - 1), claz); + if (problemElement != null) return (T)problemElement; + + problemElement = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), claz); + if (problemElement != null) return (T)problemElement; + } + return null; + } + } diff --git a/python/testData/intentions/afterTypeAssertion2.py b/python/testData/intentions/afterTypeAssertion2.py new file mode 100644 index 000000000000..eccd6b600b47 --- /dev/null +++ b/python/testData/intentions/afterTypeAssertion2.py @@ -0,0 +1,3 @@ +def foo(a, b): + assert isinstance(a, object) + a. diff --git a/python/testData/intentions/beforeTypeAssertion2.py b/python/testData/intentions/beforeTypeAssertion2.py new file mode 100644 index 000000000000..f44a3da9d087 --- /dev/null +++ b/python/testData/intentions/beforeTypeAssertion2.py @@ -0,0 +1 @@ +def foo(a, b): a. diff --git a/python/testSrc/com/jetbrains/python/PyIntentionTest.java b/python/testSrc/com/jetbrains/python/PyIntentionTest.java index cb69abafe97f..3ef18d015a58 100644 --- a/python/testSrc/com/jetbrains/python/PyIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/PyIntentionTest.java @@ -253,6 +253,10 @@ public class PyIntentionTest extends PyTestCase { doTest(PyBundle.message("INTN.insert.assertion")); } + public void testTypeAssertion2() { + doTest(PyBundle.message("INTN.insert.assertion")); + } + public void testDocStub() { doDocStubTest(); }