Create intention to suppress Mypy errors

gradle.properties

@@ -15,7 +15,7 @@
 #
 version=0.11.2
 ideVersion=PC-2021.2.3
-pythonPlugin=PythonCore:212.5457.59
+pythonPlugin=python-ce
 sinceBuild=181.5684
 untilBuild=
 downloadIdeaSources=true

src/main/java/com/leinardi/pycharm/mypy/checker/Problem.java

@@ -21,6 +21,7 @@
 import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.psi.PsiElement;
 import com.leinardi.pycharm.mypy.MypyBundle;
+import com.leinardi.pycharm.mypy.intentions.TypeIgnoreIntention;
 import com.leinardi.pycharm.mypy.mpapi.SeverityLevel;
 import org.apache.commons.lang.builder.EqualsBuilder;
 import org.apache.commons.lang.builder.HashCodeBuilder;
@@ -56,7 +57,8 @@ public void createAnnotation(@NotNull AnnotationHolder holder, @NotNull Highligh
         String message = MypyBundle.message("inspection.message", getMessage());
         AnnotationBuilder annotation = holder
                 .newAnnotation(severity, message)
-                .range(target.getTextRange());
+                .range(target.getTextRange())
+                .withFix(new TypeIgnoreIntention());
         if (isAfterEndOfLine()) {
             annotation = annotation.afterEndOfLine();
         }

src/main/java/com/leinardi/pycharm/mypy/intentions/TypeIgnoreIntention.java

@@ -0,0 +1,144 @@
+package com.leinardi.pycharm.mypy.intentions;
+
+import com.intellij.codeInsight.intention.IntentionAction;
+import com.intellij.codeInsight.intention.PsiElementBaseIntentionAction;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiComment;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.psi.LanguageLevel;
+import com.jetbrains.python.psi.PyElementGenerator;
+import com.leinardi.pycharm.mypy.MypyBundle;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Intention action to append `# type: ignore` comment to suppress Mypy annotations.
+ */
+public class TypeIgnoreIntention extends PsiElementBaseIntentionAction implements IntentionAction {
+    @NotNull
+    public String getText() {
+        return MypyBundle.message("intention.type-ignore.text");
+    }
+
+    /**
+     * This string is also used for the directory name containing the intention description.
+     */
+    @NotNull
+    public String getFamilyName() {
+        return "TypeIgnoreIntention";
+    }
+
+    @NotNull
+    String getCommentText() {
+        return "# type: ignore";
+    }
+
+    /**
+     * Checks whether this intention is available at the caret offset in file.
+     */
+    public boolean isAvailable(@NotNull Project project, Editor editor, @Nullable PsiElement element) {
+        if (element == null) {
+            return false;
+        }
+
+        PsiElement lastNode = findElementBeforeNewline(element);
+        if (lastNode == null) {
+            return false;
+        }
+
+        if (!isComment(lastNode)) {
+            // No comment - we can add one. Make sure it has a parent.
+            return lastNode.getParent() != null;
+        } else {
+            PsiComment oldComment = (PsiComment) lastNode;
+            // Extract the first part of comment, e.g.
+            // "# type: ignore # Bla bla" -> "# type: ignore"
+            String firstCommentPart = oldComment.getText().split("(?<!^)#", 2)[0].trim();
+            return !firstCommentPart.equals(getCommentText());
+        }
+    }
+
+    /**
+     * Modifies the Psi to append the `# type: ignore` comment.
+     *
+     * @throws IncorrectOperationException Thrown by underlying (Psi model) write action context
+     *                                     when manipulation of the psi tree fails.
+     */
+    public void invoke(@NotNull Project project, Editor editor, @NotNull PsiElement element)
+            throws IncorrectOperationException {
+
+        PsiElement lastNode = findElementBeforeNewline(element);
+        assert lastNode != null : "Unexpected null node";
+
+        if (isComment(lastNode)) {
+            PsiComment oldComment = (PsiComment) lastNode;
+            // Prepend to existing comment
+            String text = getCommentText() + " " + oldComment.getText();
+            oldComment.replace(createComment(element, text));
+        } else {
+            // Create a new comment at end of line
+            PsiComment newComment = createComment(lastNode, getCommentText());
+
+            assert lastNode.getParent() != null : "Unexpected null parent for " + lastNode;
+
+            // Inserting elements in the other order causes the comment to appear on the wrong line (?!)
+            PsiElement addedElement = lastNode.getParent().addAfter(newComment, lastNode);
+            addedElement.getParent().addBefore(createSpace(addedElement), addedElement);
+        }
+    }
+
+    public boolean startInWriteAction() {
+        return true;
+    }
+
+    boolean isComment(@NotNull PsiElement element) {
+        return element instanceof PsiComment
+                && ((PsiComment) element).getTokenType() == PyTokenTypes.END_OF_LINE_COMMENT;
+    }
+
+    /**
+     * Inspired by PyPsiUtils.findSameLineComment() - but that function does not behave correctly when caret is placed
+     * at the end of a line.
+     */
+    @Nullable
+    PsiElement findElementBeforeNewline(@NotNull PsiElement element) {
+        PsiElement elem = PsiTreeUtil.prevLeaf(element);
+        PsiElement next = PsiTreeUtil.getDeepestFirst(element);
+        while (true) {
+            if (next == null) {
+                // End of file
+                return elem;
+            }
+            if (next.textContains('\n')) {
+                if (next instanceof PsiWhiteSpace) {
+                    return elem;
+                } else {
+                    // If newline occurs not in whitespace (e.g. multiline string), just disable inspection for now.
+                    return null;
+                }
+            }
+            elem = next;
+            next = PsiTreeUtil.nextLeaf(next);
+        }
+    }
+
+    PsiComment createComment(@NotNull PsiElement baseElement, @NotNull String text) {
+        PyElementGenerator generator = PyElementGenerator.getInstance(baseElement.getProject());
+        return generator.createFromText(LanguageLevel.forElement(baseElement), PsiComment.class, text);
+    }
+
+    /**
+     * Generate two spaces before the comment.
+     * Per PEP 8, inline comments should be sparated by at least two spaces:
+     * https://www.python.org/dev/peps/pep-0008/#inline-comments
+     */
+    PsiWhiteSpace createSpace(@NotNull PsiElement baseElement) {
+        PyElementGenerator generator = PyElementGenerator.getInstance(baseElement.getProject());
+        return generator.createFromText(LanguageLevel.forElement(baseElement), PsiWhiteSpace.class, "  ");
+    }
+}

src/main/resources/com/leinardi/pycharm/mypy/MypyBundle.properties

@@ -61,6 +61,7 @@ plugin.exception=Unexpected Exception Caught
 inspection.group=Mypy
 inspection.display-name=Mypy real-time scan
 inspection.message=Mypy: {0}
+intention.type-ignore.text=Suppress Mypy error
 mypy.file-io-failed=<html>The Mypy rules file could not be read</html>
 mypy.exception=<html>The scan failed due to an exception: {0}</html>
 mypy.exception-with-root-cause=<html><b>The scan failed due to an exception:<b> {0}<br><b>Root cause:</b><br>{1}</html>

src/main/resources/intentionDescriptions/TypeIgnoreIntention/description.html

@@ -0,0 +1,6 @@
+<html lang="en">
+<body>
+This intention suppresses Mypy type checker errors for the given line by appending a <code># type: ignore</code>
+comment.
+</body>
+</html>

src/test/java/com/leinardi/pycharm/mypy/intentions/TypeIgnoreIntentionTest.java

@@ -0,0 +1,92 @@
+package com.leinardi.pycharm.mypy.intentions;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.testFramework.fixtures.BasePlatformTestCase;
+import com.jetbrains.python.PythonLanguage;
+
+public class TypeIgnoreIntentionTest extends BasePlatformTestCase {
+    private void assertFixerReplacement(String input, int caretPosition, String output) {
+        PsiFile file = createLightFile("test.py", PythonLanguage.INSTANCE, input);
+        PsiElement el = file.findElementAt(caretPosition);
+        assertNotNull(el);
+
+        TypeIgnoreIntention fixer = new TypeIgnoreIntention();
+        assertTrue(fixer.isAvailable(file.getProject(), null, el));
+
+        fixer.invoke(getProject(), null, el);
+        assertEquals(output, file.getText());
+    }
+
+    @SuppressWarnings("SameParameterValue")
+    private void assertNotAvailable(String input, int caretPosition) {
+        PsiFile file = createLightFile("test.py", PythonLanguage.INSTANCE, input);
+        PsiElement el = file.findElementAt(caretPosition);
+        assertNotNull(el);
+
+        TypeIgnoreIntention fixer = new TypeIgnoreIntention();
+        assertFalse(fixer.isAvailable(file.getProject(), null, el));
+    }
+
+    public void testVariableAssignment() {
+        String input = "import os\n" +
+                "\n" +
+                "foo: str = 123\n" +
+                "print()\n";
+        String output = "import os\n" +
+                "\n" +
+                "foo: str = 123  # type: ignore\n" +
+                "print()\n";
+        assertFixerReplacement(input, input.indexOf("123"), output);
+    }
+
+    public void testMultilineCall() {
+        String input = "print(\n" +
+                "    foo,\n" +
+                "    bar,\n" +
+                ")\n";
+        String output = "print(\n" +
+                "    foo,\n" +
+                "    bar,  # type: ignore\n" +
+                ")\n";
+        assertFixerReplacement(input, input.indexOf("bar"), output);
+    }
+
+    public void testMultilineString() {
+        // Trying to append comment to multiline string line would cause a syntax error.
+        String input = "foo: int = \"\"\"\n" +
+                "\"\"\"\n";
+        assertNotAvailable(input, 0);
+    }
+
+    public void testPrependComment() {
+        String input = "print()  # Hi I'm a comment\n";
+        String output = "print()  # type: ignore # Hi I'm a comment\n";
+        assertFixerReplacement(input, 0, output);
+    }
+
+    public void testCaretAtEndOfLine() {
+        String input = "print()\n\n";
+        String output = "print()  # type: ignore\n\n";
+        assertFixerReplacement(input, input.indexOf("\n"), output);
+    }
+
+    public void testNotAvailable() {
+        // `type: ignore` comment already present
+        String input = "print()  # type: ignore    #Hi I'm a comment\n";
+        assertNotAvailable(input, 0);
+    }
+
+    public void testAlmostEmptyFile() {
+        String input = "\n";
+        assertNotAvailable(input, 0);
+    }
+
+    public void testEmptyFile() {
+        String input = "";
+        // Can't use assertNotAvailable(): findElementAt() would fail.
+        PsiFile file = createLightFile("test.py", PythonLanguage.INSTANCE, input);
+        TypeIgnoreIntention fixer = new TypeIgnoreIntention();
+        assertFalse(fixer.isAvailable(file.getProject(), null, file));
+    }
+}