diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java index 098f2aba62e5..e94104501c6b 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateTypesIntention.java @@ -18,6 +18,7 @@ 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; @@ -54,7 +55,8 @@ 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 PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + 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); @@ -77,7 +79,8 @@ public class PyAnnotateTypesIntention extends PyBaseIntentionAction { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + 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) { diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java index 071bb341ca4a..5658e7276b87 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInDocstringIntention.java @@ -16,6 +16,7 @@ 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; @@ -57,7 +58,8 @@ public class SpecifyTypeInDocstringIntention extends TypeIntention { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + 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(); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java index fb9834921220..69fc69bbfa9e 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/SpecifyTypeInPy3AnnotationsIntention.java @@ -17,6 +17,7 @@ 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; @@ -68,7 +69,8 @@ public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention { @Override public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + 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(); diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java index f6ed28246438..17d48e4d7ef8 100644 --- a/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java +++ b/python/src/com/jetbrains/python/codeInsight/intentions/TypeIntention.java @@ -15,6 +15,7 @@ */ package com.jetbrains.python.codeInsight.intentions; +import com.intellij.codeInsight.TargetElementUtil; import com.intellij.lang.ASTNode; import com.intellij.openapi.editor.Editor; import com.intellij.openapi.project.Project; @@ -24,7 +25,6 @@ import com.intellij.openapi.vfs.VirtualFile; import com.intellij.psi.*; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.util.containers.ContainerUtil; -import com.jetbrains.python.PyNames; import com.jetbrains.python.documentation.doctest.PyDocstringFile; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.resolve.PyResolveContext; @@ -49,7 +49,8 @@ public abstract class TypeIntention extends PyBaseIntentionAction { if (!(file instanceof PyFile) || file instanceof PyDocstringFile) return false; updateText(false); - final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset()); + 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)) { return true; @@ -86,15 +87,7 @@ public abstract class TypeIntention extends PyBaseIntentionAction { @Nullable public static PyExpression getProblemElement(@Nullable PsiElement elementAt) { - PyExpression problemElement = PsiTreeUtil.getParentOfType(elementAt, PyNamedParameter.class, PyReferenceExpression.class); - if (problemElement == null) return null; - if (problemElement instanceof PyQualifiedExpression) { - final PyExpression qualifier = ((PyQualifiedExpression)problemElement).getQualifier(); - if (qualifier != null && !qualifier.getText().equals(PyNames.CANONICAL_SELF)) { - problemElement = qualifier; - } - } - return problemElement; + return PsiTreeUtil.getParentOfType(elementAt, PyNamedParameter.class, PyReferenceExpression.class); } protected abstract void updateText(boolean isReturn); diff --git a/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall.py b/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall.py new file mode 100644 index 000000000000..0cdc87e27827 --- /dev/null +++ b/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall.py @@ -0,0 +1,10 @@ +class MyClass: + def __init__(self): + pass + + def method(self, x): + pass + + +x = MyClass() +foo = x.method(42) diff --git a/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall_after.py b/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall_after.py new file mode 100644 index 000000000000..619d2671d5f4 --- /dev/null +++ b/python/testData/intentions/PyAnnotateTypesIntentionTest/methodAfterConstructorCall_after.py @@ -0,0 +1,11 @@ +class MyClass: + def __init__(self): + pass + + def method(self, x): + # type: (object) -> object + pass + + +x = MyClass() +foo = x.method(42) diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java index 364e79895e23..95610a7205d6 100644 --- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateTypesIntentionTest.java @@ -46,6 +46,11 @@ public class PyAnnotateTypesIntentionTest extends PyIntentionTestCase { doNegativeTest(PyBundle.message("INTN.add.type.hints.for.function.family")); } + // PY-30825 + public void testMethodAfterConstructorCall() { + doIntentionTest(PyBundle.message("INTN.add.type.hints.for.function", "method")); + } + private void doTest() { doTest(PyBundle.message("INTN.add.type.hints.for.function.family"), LanguageLevel.PYTHON34); }