Auto import actions using importmagic

youben11 · youben11 · commit e7e7dd25ceb3 · 2019-10-28T22:16:03.000+01:00
diff --git a/pyls/_utils.py b/pyls/_utils.py
@@ -192,3 +192,7 @@ def is_process_alive(pid):
             return e.errno == errno.EPERM
         else:
             return True
+
+
+def camel_to_underscore(camelcase):
+    return re.sub('([A-Z]+)', r'_\1',camelcase).lower()
diff --git a/pyls/plugins/importmagic_lint.py b/pyls/plugins/importmagic_lint.py
@@ -0,0 +1,180 @@
+# Copyright 2017 Palantir Technologies, Inc.
+import logging
+import re
+import sys
+import importmagic
+from pyls import hookimpl, lsp, _utils
+
+
+log = logging.getLogger(__name__)
+
+SOURCE = 'importmagic'
+ADD_IMPORT_COMMAND = 'importmagic.addimport'
+MAX_COMMANDS = 4
+UNRES_RE = re.compile(r"Unresolved import '(?P<unresolved>[\w.]+)'")
+
+_index_cache = {}
+
+
+def _get_index(sys_path):
+    """Build index of symbols from python modules.
+    Cache the index so we don't build it multiple times unnecessarily.
+    """
+    key = tuple(sys_path)
+    if key not in _index_cache:
+        log.debug("Started building importmagic index")
+        index = importmagic.SymbolIndex()
+        # The build tend to be noisy
+        logging.getLogger('importmagic.index').setLevel(logging.ERROR)
+        index.build_index(paths=sys_path)
+        _index_cache[key] = index
+        logging.getLogger('importmagic.index').setLevel(logging.DEBUG)
+        log.debug("Finished building importmagic index")
+    return _index_cache[key]
+
+
+@hookimpl
+def pyls_commands():
+    return [ADD_IMPORT_COMMAND]
+
+
+@hookimpl
+def pyls_lint(document):
+    """Build a diagnostics of unresolved symbols. Every entry follows this format:
+        {
+            'source': 'importmagic',
+            'range': {
+                'start': {
+                    'line': start_line,
+                    'character': start_column,
+                },
+                'end': {
+                    'line': end_line,
+                    'character': end_column,
+                },
+            },
+            'message': 'Unresolved import <symbol>',
+            'severity': lsp.DiagnosticSeverity.Hint,
+        }
+
+    Args:
+        document: The document to be linted.
+    Returns:
+        A list of dictionaries.
+    """
+    scope = importmagic.Scope.from_source(document.source)
+    unresolved, _unreferenced = scope.find_unresolved_and_unreferenced_symbols()
+
+    diagnostics = []
+
+    # Annoyingly, we only get the text of an unresolved import, so we'll look for it ourselves
+    for unres in unresolved:
+        if unres not in document.source:
+            continue
+
+        for line_no, line in enumerate(document.lines):
+            pos = line.find(unres)
+            if pos < 0:
+                continue
+
+            diagnostics.append({
+                'source': SOURCE,
+                'range': {
+                    'start': {'line': line_no, 'character': pos},
+                    'end': {'line': line_no, 'character': pos + len(unres)}
+                },
+                'message': "Unresolved import '%s'" % unres,
+                'severity': lsp.DiagnosticSeverity.Hint,
+            })
+
+    return diagnostics
+
+
+@hookimpl
+def pyls_code_actions(config, document, context):
+    """Build a list of actions to be suggested to the user. Each action follow this format:
+        {
+            'title': 'importmagic',
+            'command': command ('importmagic.add_import'),
+            'arguments':
+                {
+                    'uri': document.uri,
+                    'version': document.version,
+                    'startLine': start_line,
+                    'endLine': end_line,
+                    'newText': text,
+                }
+        }
+    """
+    # Update the style configuration
+    conf = config.plugin_settings('importmagic_lint')
+    log.debug("Got importmagic settings: %s", conf)
+    importmagic.Imports.set_style(**{_utils.camel_to_underscore(k): v for k, v in conf.items()})
+
+    actions = []
+    diagnostics = context.get('diagnostics', [])
+    for diagnostic in diagnostics:
+        if diagnostic.get('source') != SOURCE:
+            continue
+        m = UNRES_RE.match(diagnostic['message'])
+        if not m:
+            continue
+
+        unres = m.group('unresolved')
+        # Might be slow but is cached once built
+        index = _get_index(sys.path)
+
+        for score, module, variable in sorted(index.symbol_scores(unres)[:MAX_COMMANDS], reverse=True):
+            if score < 1:
+                # These tend to be terrible
+                continue
+
+            # Generate the patch we would need to apply
+            imports = importmagic.Imports(index, document.source)
+            if variable:
+                imports.add_import_from(module, variable)
+            else:
+                imports.add_import(module)
+            start_line, end_line, text = imports.get_update()
+
+            actions.append({
+                'title': _command_title(variable, module),
+                'command': ADD_IMPORT_COMMAND,
+                'arguments': [{
+                    'uri': document.uri,
+                    'version': document.version,
+                    'startLine': start_line,
+                    'endLine': end_line,
+                    'newText': text
+                }]
+            })
+    return actions
+
+
+@hookimpl
+def pyls_execute_command(workspace, command, arguments):
+    if command != ADD_IMPORT_COMMAND:
+        return
+
+    args = arguments[0]
+
+    edit = {'documentChanges': [{
+        'textDocument': {
+            'uri': args['uri'],
+            'version': args['version']
+        },
+        'edits': [{
+            'range': {
+                'start': {'line': args['startLine'], 'character': 0},
+                'end': {'line': args['endLine'], 'character': 0},
+            },
+            'newText': args['newText']
+        }]
+    }]}
+    workspace.apply_edit(edit)
+
+
+def _command_title(variable, module):
+    if not variable:
+        return 'Import "%s"' % module
+    return 'Import "%s" from "%s"' % (variable, module)
diff --git a/setup.py b/setup.py
@@ -48,6 +48,7 @@
         'all': [
             'autopep8',
             'flake8',
+            'importmagic',
             'mccabe',
             'pycodestyle',
             'pydocstyle>=2.0.0',
@@ -58,6 +59,7 @@
         ],
         'autopep8': ['autopep8'],
         'flake8': ['flake8'],
+        'importmagic': ['importmagic'],
         'mccabe': ['mccabe'],
         'pycodestyle': ['pycodestyle'],
         'pydocstyle': ['pydocstyle>=2.0.0'],
@@ -78,6 +80,7 @@
         'pyls': [
             'autopep8 = pyls.plugins.autopep8_format',
             'flake8 = pyls.plugins.flake8_lint',
+            'importmagic = pyls.plugins.importmagic_lint',
             'jedi_completion = pyls.plugins.jedi_completion',
             'jedi_definition = pyls.plugins.definition',
             'jedi_hover = pyls.plugins.hover',
diff --git a/test/plugins/test_importmagic_lint.py b/test/plugins/test_importmagic_lint.py
@@ -0,0 +1,68 @@
+# Copyright 2019 Palantir Technologies, Inc.
+import tempfile
+import os
+from pyls import lsp, uris
+from pyls.plugins import importmagic_lint
+from pyls.workspace import Document
+
+DOC_URI = uris.from_fs_path(__file__)
+DOC = """
+files = listdir()
+print(files)
+"""
+
+
+def temp_document(doc_text):
+    temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
+    name = temp_file.name
+    temp_file.write(doc_text)
+    temp_file.close()
+    doc = Document(uris.from_fs_path(name))
+
+    return name, doc
+
+
+def test_importmagic_lint():
+    try:
+        name, doc = temp_document(DOC)
+        diags = importmagic_lint.pyls_lint(doc)
+        unres_symbol = [d for d in diags if d['source'] == 'importmagic'][0]
+
+        assert unres_symbol['message'] == "Unresolved import 'listdir'"
+        assert unres_symbol['range']['start'] == {'line': 1, 'character': 8}
+        assert unres_symbol['range']['end'] == {'line': 1, 'character': 15}
+        assert unres_symbol['severity'] == lsp.DiagnosticSeverity.Hint
+
+    finally:
+        os.remove(name)
+
+
+def test_importmagic_actions(config):
+    context = {
+        'diagnostict': [
+            {
+                'range':
+                {
+                    'start': {'line': 1, 'character': 8},
+                    'end': {'line': 1, 'character': 15}
+                },
+                'message': "Unresolved import 'listdir'",
+                'severity': 4,
+                'source': 'importmagic'
+            }
+        ]
+    }
+
+    try:
+        name, doc = temp_document(DOC)
+        action = importmagic_lint.pyls_code_actions(config, doc, context)[0]
+        arguments = action['arguments']
+
+        assert action['title'] == 'Import "listdir" from "os"'
+        assert action['command'] == 'importmagic.addimport'
+        assert arguments['startLine'] == 1
+        assert arguments['endLine'] == 1
+        assert arguments['newText'] == 'from os import listdir\n\n\n'
+
+    finally:
+        os.remove(name)
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -87,3 +87,9 @@ def test_clip_column():
     assert _utils.clip_column(2, ['123\n', '123'], 0) == 2
     assert _utils.clip_column(3, ['123\n', '123'], 0) == 3
     assert _utils.clip_column(4, ['123\n', '123'], 1) == 3
+
+
+def test_camel_to_underscore():
+    assert _utils.camel_to_underscore('camelCase') == 'camel_case'
+    assert _utils.camel_to_underscore('hangClosing') == 'hang_closing'
+    assert _utils.camel_to_underscore('ignore') == 'ignore'
