non blocking index build

Author: youben11

pyls/plugins/importmagic_lint.py

@@ -2,6 +2,7 @@
 import logging
 import re
 import sys
+from concurrent.futures import ThreadPoolExecutor
 import importmagic
 from pyls import hookimpl, lsp, _utils
 
@@ -15,22 +16,37 @@
 UNRES_RE = re.compile(r"Unresolved import '(?P<unresolved>[\w.]+)'")
 UNREF_RE = re.compile(r"Unreferenced import '(?P<unreferenced>[\w.]+)'")
 
-_index_cache = {}
+_index_cache = None
 
 
-def _get_index(sys_path):
+def _build_index(paths):
     """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.info("Started building importmagic index")
-        index = importmagic.SymbolIndex()
-        # The build tend to be noisy
-        index.build_index(paths=sys_path)
-        _index_cache[key] = index
-        log.info("Finished building importmagic index")
-    return _index_cache[key]
+    log.info("Started building importmagic index")
+    index = importmagic.SymbolIndex()
+    index.build_index(paths=paths)
+    log.info("Finished building importmagic index")
+    return index
+
+
+def _cache_index_callback(future):
+    global _index_cache
+    # Cache the index
+    _index_cache = future.result()
+
+
+def _get_index():
+    """Get the cached index if built and index project files on each call.
+    Return an empty index if not built yet.
+    """
+    # Index haven't been built yet
+    if _index_cache is None:
+        return importmagic.SymbolIndex()
+
+    # Index project files
+    # TODO(youben) index project files
+    #index.build_index(paths=[])
+    return _index_cache
 
 
 def _get_imports_list(source, index=None):
@@ -46,6 +62,13 @@ def _get_imports_list(source, index=None):
     return imported
 
 
+@hookimpl
+def pyls_initialize():
+    pool = ThreadPoolExecutor()
+    builder = pool.submit(_build_index, (sys.path))
+    builder.add_done_callback(_cache_index_callback)
+
+
 @hookimpl
 def pyls_commands():
     return [ADD_IMPORT_COMMAND, REMOVE_IMPORT_COMMAND]
@@ -125,7 +148,7 @@ def pyls_lint(document):
 
 
 @hookimpl
-def pyls_code_actions(config, document):
+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',
@@ -146,9 +169,8 @@ def pyls_code_actions(config, document):
     log.debug("Got importmagic settings: %s", conf)
     importmagic.Imports.set_style(**{_utils.camel_to_underscore(k): v for k, v in conf.items()})
 
-    # Might be slow but is cached once built
-    # TODO (youben): add project path for indexing
-    index = _get_index(sys.path)
+    # Get empty index while it's building so we don't block here
+    index = _get_index()
     actions = []
 
     diagnostics = pyls_lint(document)

test/plugins/test_importmagic_lint.py

@@ -1,6 +1,7 @@
 # Copyright 2019 Palantir Technologies, Inc.
 import tempfile
 import os
+from time import sleep
 from pyls import lsp, uris
 from pyls.plugins import importmagic_lint
 from pyls.workspace import Document
@@ -39,7 +40,11 @@ def test_importmagic_lint():
 
 def test_importmagic_actions(config):
     try:
+        importmagic_lint.pyls_initialize()
         name, doc = temp_document(DOC)
+        while importmagic_lint._index_cache is None:
+            # wait for the index to be ready
+            sleep(1)
         actions = importmagic_lint.pyls_code_actions(config, doc)
         action = [a for a in actions if a['title'] == 'Import "time"'][0]
         arguments = action['arguments'][0]
@@ -51,3 +56,5 @@ def test_importmagic_actions(config):
 
     finally:
         os.remove(name)
+
+# TODO(youben) write test for remove action