Add code actions for simple Pylint items (#300)

Co-authored-by: Karthik Nadig <[email protected]>

Author: dciborow

bundled/tool/lsp_server.py

@@ -7,6 +7,7 @@
 import json
 import os
 import pathlib
+import re
 import sys
 import traceback
 from typing import Any, Callable, Dict, List, Optional, Sequence, Union
@@ -311,6 +312,60 @@ def organize_imports(
     ]
 
 
+REPLACEMENTS = {
+    "R0205:useless-object-inheritance": {
+        "pattern": r"class (\w+)\(object\):",
+        "repl": r"class \1:",
+    },
+    "R1721:unnecessary-comprehension": {
+        "pattern": r"\{([\w\s,]+) for [\w\s,]+ in ([\w\s,]+)\}",
+        "repl": r"set(\2)",
+    },
+    "E1141:dict-iter-missing-items": {
+        "pattern": r"for\s+(\w+),\s+(\w+)\s+in\s+(\w+)\s*:",
+        "repl": r"for \1, \2 in \3.items():",
+    },
+}
+
+
+def _get_replacement_edit(diagnostic: lsp.Diagnostic, lines: List[str]) -> lsp.TextEdit:
+    return lsp.TextEdit(
+        lsp.Range(
+            start=lsp.Position(line=diagnostic.range.start.line, character=0),
+            end=lsp.Position(line=diagnostic.range.start.line + 1, character=0),
+        ),
+        re.sub(
+            REPLACEMENTS[diagnostic.code]["pattern"],
+            REPLACEMENTS[diagnostic.code]["repl"],
+            lines[diagnostic.range.start.line],
+        ),
+    )
+
+
+@QUICK_FIXES.quick_fix(
+    codes=list(REPLACEMENTS.keys()),
+)
+def fix_with_replacement(
+    document: workspace.Document, diagnostics: List[lsp.Diagnostic]
+) -> List[lsp.CodeAction]:
+    """Provides quick fixes which basic string replacements."""
+    return [
+        lsp.CodeAction(
+            title=f"{TOOL_DISPLAY}: Run autofix code action",
+            kind=lsp.CodeActionKind.QuickFix,
+            diagnostics=diagnostics,
+            edit=_create_workspace_edits(
+                document,
+                [
+                    _get_replacement_edit(diagnostic, document.lines)
+                    for diagnostic in diagnostics
+                    if diagnostic.code in REPLACEMENTS
+                ],
+            ),
+        )
+    ]
+
+
 def _command_quick_fix(
     diagnostics: List[lsp.Diagnostic],
     title: str,
@@ -333,7 +388,7 @@ def _create_workspace_edits(
             lsp.TextDocumentEdit(
                 text_document=lsp.OptionalVersionedTextDocumentIdentifier(
                     uri=document.uri,
-                    version=document.version,
+                    version=document.version if document.version else 0,
                 ),
                 edits=results,
             )

src/test/python_tests/test_code_actions.py

@@ -137,3 +137,113 @@ def _handler(params):
             ]
 
         assert_that(actual_code_actions, is_(expected))
+
+
+@pytest.mark.parametrize(
+    ("code", "contents", "new_text"),
+    [
+        (
+            "R0205:useless-object-inheritance",
+            """
+class Banana(object):
+    pass""",
+            """class Banana:
+""",
+        ),
+        (
+            "R1721:unnecessary-comprehension",
+            """
+NUMBERS = [1, 1, 2, 2, 3, 3]
+
+UNIQUE_NUMBERS = {number for number in NUMBERS}
+""",
+            """UNIQUE_NUMBERS = set(NUMBERS)
+""",
+        ),
+        (
+            "E1141:dict-iter-missing-items",
+            """
+data = {'Paris': 2_165_423, 'New York City': 8_804_190, 'Tokyo': 13_988_129}
+for city, population in data:
+    print(f"{city} has population {population}.")
+""",
+            """for city, population in data.items():
+""",
+        ),
+    ],
+)
+def test_edit_code_action(code, contents, new_text):
+    """Tests for code actions which run a command."""
+    with utils.python_file(contents, TEST_FILE_PATH.parent) as temp_file:
+        uri = utils.as_uri(os.fspath(temp_file))
+
+        actual = {}
+        with session.LspSession() as ls_session:
+            ls_session.initialize()
+
+            done = Event()
+
+            def _handler(params):
+                nonlocal actual
+                actual = params
+                done.set()
+
+            ls_session.set_notification_callback(session.PUBLISH_DIAGNOSTICS, _handler)
+
+            ls_session.notify_did_open(
+                {
+                    "textDocument": {
+                        "uri": uri,
+                        "languageId": "python",
+                        "version": 1,
+                        "text": contents,
+                    }
+                }
+            )
+
+            # wait for some time to receive all notifications
+            done.wait(TIMEOUT)
+
+            diagnostics = [d for d in actual["diagnostics"] if d["code"] == code]
+
+            assert_that(len(diagnostics), is_(greater_than(0)))
+
+            actual_code_actions = ls_session.text_document_code_action(
+                {
+                    "textDocument": {"uri": uri},
+                    "range": {
+                        "start": {"line": 0, "character": 0},
+                        "end": {"line": 1, "character": 0},
+                    },
+                    "context": {"diagnostics": diagnostics},
+                }
+            )
+            text_document = actual_code_actions[0]["edit"]["documentChanges"][0][
+                "textDocument"
+            ]
+            text_range = actual_code_actions[0]["edit"]["documentChanges"][0]["edits"][
+                0
+            ]["range"]
+            expected = [
+                {
+                    "title": f"{LINTER}: Run autofix code action",
+                    "kind": "quickfix",
+                    "diagnostics": [d],
+                    "edit": {
+                        "documentChanges": [
+                            {
+                                "textDocument": text_document,
+                                "edits": [
+                                    {
+                                        "range": text_range,
+                                        "newText": new_text,
+                                    }
+                                ],
+                            }
+                        ]
+                    },
+                }
+                for d in diagnostics
+            ]
+
+        assert_that(actual_code_actions, is_(expected))