diff --git a/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/after.py.template b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/after.py.template new file mode 100644 index 000000000000..37426b41ba70 --- /dev/null +++ b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/after.py.template @@ -0,0 +1,5 @@ +class C: + attr: str + + def __init__(self, attr='spam'): + self.attr = attr diff --git a/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/before.py.template b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/before.py.template new file mode 100644 index 000000000000..de1a13828690 --- /dev/null +++ b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/before.py.template @@ -0,0 +1,3 @@ +class C: + def __init__(self, attr='spam'): + self.attr = attr diff --git a/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/description.html b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/description.html new file mode 100644 index 000000000000..01de4a5112fe --- /dev/null +++ b/python/resources/intentionDescriptions/PyAnnotateVariableTypeIntention/description.html @@ -0,0 +1,6 @@ + + +This intention adds type hints for variables in PEP 484 compatible format. + + + \ No newline at end of file diff --git a/python/src/META-INF/python-core-common.xml b/python/src/META-INF/python-core-common.xml index 3fde43ea74f9..7f4dd155634a 100644 --- a/python/src/META-INF/python-core-common.xml +++ b/python/src/META-INF/python-core-common.xml @@ -340,6 +340,11 @@ Python + + com.jetbrains.python.codeInsight.intentions.PyAnnotateVariableTypeIntention + Python + + diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java new file mode 100644 index 000000000000..05248c5ec2fd --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java @@ -0,0 +1,92 @@ +// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +package com.jetbrains.python.codeInsight.intentions; + +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.project.Project; +import com.intellij.openapi.roots.ProjectFileIndex; +import com.intellij.psi.PsiFile; +import com.intellij.psi.util.PsiTreeUtil; +import com.intellij.util.IncorrectOperationException; +import com.jetbrains.python.PyBundle; +import com.jetbrains.python.codeInsight.typing.PyTypingTypeProvider; +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.TypeEvalContext; +import one.util.streamex.StreamEx; +import org.jetbrains.annotations.Nls; +import org.jetbrains.annotations.NotNull; + +import java.util.Collections; +import java.util.List; + +/** + * @author Mikhail Golubev + */ +public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction { + @Nls + @NotNull + @Override + public String getFamilyName() { + return PyBundle.message("INTN.annotate.types"); + } + + @Override + public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) { + if (!(file instanceof PyFile) || file instanceof PyDocstringFile) { + return false; + } + final List resolved = findSuitableTargetsUnderCaret(project, editor, file); + if (resolved.isEmpty() || resolved.size() > 1) { + return false; + } + + setText(PyBundle.message("INTN.annotate.types")); + return true; + } + + @NotNull + private static List findSuitableTargetsUnderCaret(@NotNull Project project, Editor editor, PsiFile file) { + final PyReferenceOwner elementAtCaret = PsiTreeUtil.getParentOfType(file.findElementAt(editor.getCaretModel().getOffset()), + PyReferenceExpression.class, PyTargetExpression.class); + if (elementAtCaret == null) { + return Collections.emptyList(); + } + + final ProjectFileIndex index = ProjectFileIndex.getInstance(project); + final TypeEvalContext typeEvalContext = TypeEvalContext.codeAnalysis(project, file); + final PyResolveContext resolveContext = PyResolveContext.noImplicits().withTypeEvalContext(typeEvalContext); + // TODO filter out targets defined in stubs + return StreamEx.of(PyUtil.multiResolveTopPriority(elementAtCaret, resolveContext)) + .select(PyTargetExpression.class) + .filter(target -> !index.isInLibraryClasses(target.getContainingFile().getVirtualFile())) + .filter(target -> !isAnnotated(target, typeEvalContext)) + .toList(); + } + + private static boolean isAnnotated(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) { + // TODO filter out fields explicitly annotated as Any + return new PyTypingTypeProvider().getReferenceType(target, context, null) != null; + } + + @Override + public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException { + final List targets = findSuitableTargetsUnderCaret(project, editor, file); + assert targets.size() == 1; + final PyTargetExpression annotationTarget = targets.get(0); + if (LanguageLevel.forElement(annotationTarget).isAtLeast(LanguageLevel.PYTHON36)) { + insertVariableAnnotation(annotationTarget); + } + else { + insertVariableTypeComment(annotationTarget); + } + } + + private void insertVariableAnnotation(@NotNull PyTargetExpression target) { + + } + + private void insertVariableTypeComment(@NotNull PyTargetExpression target) { + + } +} diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget.py new file mode 100644 index 000000000000..bfdb8566d32c --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget.py @@ -0,0 +1,3 @@ +def func(): + var = 'spam' + var diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget_after.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget_after.py new file mode 100644 index 000000000000..ff144384a43d --- /dev/null +++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/annotationLocalAssignmentTarget_after.py @@ -0,0 +1,3 @@ +def func(): + var: str = 'spam' + var diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java new file mode 100644 index 000000000000..e033b5c1dd39 --- /dev/null +++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java @@ -0,0 +1,27 @@ +// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. +package com.jetbrains.python.intentions; + +import com.jetbrains.python.PyBundle; +import com.jetbrains.python.psi.LanguageLevel; +import org.jetbrains.annotations.NotNull; + +/** + * @author Mikhail Golubev + */ +public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase { + public void testAnnotationLocalAssignmentTarget() { + doTestAnnotation(); + } + + private void doTestAnnotation() { + doTest(LanguageLevel.PYTHON36); + } + + private void doTestTypeComment() { + doTest(LanguageLevel.PYTHON27); + } + + private void doTest(@NotNull LanguageLevel languageLevel) { + doTest(PyBundle.message("INTN.annotate.types"), languageLevel); + } +}