Implement cached label resolution and label resolution limit

Author: krassowski

pylsp/plugins/jedi_completion.py

@@ -3,7 +3,10 @@
 
 import logging
 import os.path as osp
+from collections import defaultdict
+from time import time
 
+from jedi.api.classes import Completion
 import parso
 
 from pylsp import _utils, hookimpl, lsp
@@ -55,6 +58,7 @@
 @hookimpl
 def pylsp_completions(config, document, position):
     """Get formatted completions for current code position"""
+    # pylint: disable=too-many-locals
     settings = config.plugin_settings('jedi_completion', document_path=document.path)
     code_position = _utils.position_to_jedi_linecolumn(document, position)
 
@@ -70,18 +74,28 @@ def pylsp_completions(config, document, position):
     should_include_params = settings.get('include_params')
     should_include_class_objects = settings.get('include_class_objects', True)
 
+    max_labels_resolve = settings.get('resolve_at_most_labels', 25)
+
     include_params = snippet_support and should_include_params and use_snippets(document, position)
     include_class_objects = snippet_support and should_include_class_objects and use_snippets(document, position)
 
     ready_completions = [
-        _format_completion(c, include_params)
-        for c in completions
+        _format_completion(
+            c,
+            include_params,
+            resolve_label=(i < max_labels_resolve)
+        )
+        for i, c in enumerate(completions)
     ]
 
     if include_class_objects:
-        for c in completions:
+        for i, c in enumerate(completions):
             if c.type == 'class':
-                completion_dict = _format_completion(c, False)
+                completion_dict = _format_completion(
+                    c,
+                    False,
+                    resolve_label=(i < max_labels_resolve)
+                )
                 completion_dict['kind'] = lsp.CompletionItemKind.TypeParameter
                 completion_dict['label'] += ' object'
                 ready_completions.append(completion_dict)
@@ -139,9 +153,9 @@ def use_snippets(document, position):
             not (expr_type in _ERRORS and 'import' in code))
 
 
-def _format_completion(d, include_params=True):
+def _format_completion(d, include_params=True, resolve_label=False):
     completion = {
-        'label': _label(d),
+        'label': _label(d, resolve_label),
         'kind': _TYPE_MAP.get(d.type),
         'detail': _detail(d),
         'documentation': _utils.format_docstring(d.docstring()),
@@ -180,12 +194,12 @@ def _format_completion(d, include_params=True):
     return completion
 
 
-def _label(definition):
-    sig = definition.get_signatures()
-    if definition.type in ('function', 'method') and sig:
-        params = ', '.join(param.name for param in sig[0].params)
-        return '{}({})'.format(definition.name, params)
-
+def _label(definition, resolve=False):
+    if not resolve:
+        return definition.name
+    sig = LABEL_RESOLVER.get_or_create(definition)
+    if sig:
+        return sig
     return definition.name
 
 
@@ -204,3 +218,78 @@ def _sort_text(definition):
     # If its 'hidden', put it next last
     prefix = 'z{}' if definition.name.startswith('_') else 'a{}'
     return prefix.format(definition.name)
+
+
+class LabelResolver:
+
+    def __init__(self, format_label_callback, time_to_live=60 * 30):
+        self.format_label = format_label_callback
+        self._cache = {}
+        self._time_to_live = time_to_live
+        self._cache_ttl = defaultdict(set)
+        self._clear_every = 2
+        # see https://github.com/davidhalter/jedi/blob/master/jedi/inference/helpers.py#L194-L202
+        self._cached_modules = {'pandas', 'numpy', 'tensorflow', 'matplotlib'}
+
+    def clear_outdated(self):
+        now = self.time_key()
+        to_clear = [
+            timestamp
+            for timestamp in self._cache_ttl
+            if timestamp < now
+        ]
+        for time_key in to_clear:
+            for key in self._cache_ttl[time_key]:
+                del self._cache[key]
+            del self._cache_ttl[time_key]
+
+    def time_key(self):
+        return int(time() / self._time_to_live)
+
+    def get_or_create(self, completion: Completion):
+        if not completion.full_name:
+            use_cache = False
+        else:
+            module_parts = completion.full_name.split('.')
+            use_cache = module_parts and module_parts[0] in self._cached_modules
+
+        if use_cache:
+            key = self._create_completion_id(completion)
+            if key not in self._cache:
+                if self.time_key() % self._clear_every == 0:
+                    self.clear_outdated()
+
+                self._cache[key] = self.resolve_label(completion)
+                self._cache_ttl[self.time_key()].add(key)
+            return self._cache[key]
+
+        return self.resolve_label(completion)
+
+    def _create_completion_id(self, completion: Completion):
+        return (
+            completion.full_name, completion.module_path,
+            completion.line, completion.column,
+            self.time_key()
+        )
+
+    def resolve_label(self, completion):
+        try:
+            sig = completion.get_signatures()
+            return self.format_label(completion, sig)
+        except Exception as e:  # pylint: disable=broad-except
+            log.warning(
+                'Something went wrong when resolving label for {completion}: {e}',
+                completion=completion, e=e
+            )
+            return ''
+
+
+def format_label(completion, sig):
+    if sig and completion.type in ('function', 'method'):
+        params = ', '.join(param.name for param in sig[0].params)
+        label = '{}({})'.format(completion.name, params)
+        return label
+    return completion.name
+
+
+LABEL_RESOLVER = LabelResolver(format_label)

test/plugins/test_completion.py

@@ -1,6 +1,7 @@
 # Copyright 2017-2020 Palantir Technologies, Inc.
 # Copyright 2021- Python Language Server Contributors.
 
+import math
 import os
 import sys
 
@@ -79,6 +80,24 @@ def test_jedi_completion_with_fuzzy_enabled(config, workspace):
     pylsp_jedi_completions(config, doc, {'line': 1, 'character': 1000})
 
 
+def test_jedi_completion_resolve_at_most(config, workspace):
+    # Over 'i' in os.path.isabs(...)
+    com_position = {'line': 1, 'character': 15}
+    doc = Document(DOC_URI, workspace, DOC)
+
+    # Do not resolve any labels
+    config.update({'plugins': {'jedi_completion': {'resolve_at_most_labels': 0}}})
+    items = pylsp_jedi_completions(config, doc, com_position)
+    labels = {i['label'] for i in items}
+    assert 'isabs' in labels
+
+    # Resolve all items
+    config.update({'plugins': {'jedi_completion': {'resolve_at_most_labels': math.inf}}})
+    items = pylsp_jedi_completions(config, doc, com_position)
+    labels = {i['label'] for i in items}
+    assert 'isabs(path)' in labels
+
+
 def test_rope_completion(config, workspace):
     # Over 'i' in os.path.isabs(...)
     com_position = {'line': 1, 'character': 15}
@@ -94,6 +113,7 @@ def test_jedi_completion_ordering(config, workspace):
     # Over the blank line
     com_position = {'line': 8, 'character': 0}
     doc = Document(DOC_URI, workspace, DOC)
+    config.update({'plugins': {'jedi_completion': {'resolve_at_most_labels': math.inf}}})
     completions = pylsp_jedi_completions(config, doc, com_position)
 
     items = {c['label']: c['sortText'] for c in completions}