Implement cached label resolution and label resolution limit (#26)

* Implement cached label resolution and label resolution limit

* Add settings to schema

* Update test for new option

* Allow to customize the modules for which the labels should be cached

* Fix the default

Author: krassowski

CONFIGURATION.md

@@ -12,6 +12,8 @@ This server can be configured using `workspace/didChangeConfiguration` method. E
 | `pylsp.plugins.jedi_completion.include_class_objects` | `boolean` | Adds class objects as a separate completion item. | `true` |
 | `pylsp.plugins.jedi_completion.fuzzy` | `boolean` | Enable fuzzy when requesting autocomplete. | `false` |
 | `pylsp.plugins.jedi_completion.eager` | `boolean` | Resolve documentation and detail eagerly. | `false` |
+| `pylsp.plugins.jedi_completion.resolve_at_most_labels` | `number`  | How many labels (at most) should be resolved? | `25` |
+| `pylsp.plugins.jedi_completion.cache_labels_for` | `array`  of  `string` items | Modules for which the labels should be cached. | `["pandas", "numpy", "tensorflow", "matplotlib"]` |
 | `pylsp.plugins.jedi_definition.enabled` | `boolean` | Enable or disable the plugin. | `true` |
 | `pylsp.plugins.jedi_definition.follow_imports` | `boolean` | The goto call will follow imports. | `true` |
 | `pylsp.plugins.jedi_definition.follow_builtin_imports` | `boolean` | If follow_imports is True will decide if it follow builtin imports. | `true` |
@@ -44,6 +46,7 @@ This server can be configured using `workspace/didChangeConfiguration` method. E
 | `pylsp.plugins.pylint.args` | `array`  of non-unique `string` items | Arguments to pass to pylint. | `null` |
 | `pylsp.plugins.pylint.executable` | `string` | Executable to run pylint with. Enabling this will run pylint on unsaved files via stdin. Can slow down workflow. Only works with python3. | `null` |
 | `pylsp.plugins.rope_completion.enabled` | `boolean` | Enable or disable the plugin. | `true` |
+| `pylsp.plugins.rope_completion.eager` | `boolean` | Resolve documentation and detail eagerly. | `false` |
 | `pylsp.plugins.yapf.enabled` | `boolean` | Enable or disable the plugin. | `true` |
 | `pylsp.rope.extensionModules` | `string` | Builtin and c-extension modules that are allowed to be imported and inspected by rope. | `null` |
 | `pylsp.rope.ropeFolder` | `array`  of unique `string` items | The name of the folder in which rope stores project configurations and data.  Pass `null` for not using such a folder at all. | `null` |

pylsp/config/schema.json

@@ -54,6 +54,19 @@
       "default": false,
       "description": "Resolve documentation and detail eagerly."
     },
+    "pylsp.plugins.jedi_completion.resolve_at_most_labels": {
+      "type": "number",
+      "default": 25,
+      "description": "How many labels (at most) should be resolved?"
+    },
+    "pylsp.plugins.jedi_completion.cache_labels_for": {
+      "type": "array",
+      "items": {
+        "type": "string"
+      },
+      "default": ["pandas", "numpy", "tensorflow", "matplotlib"],
+      "description": "Modules for which the labels should be cached."
+    },
     "pylsp.plugins.jedi_definition.enabled": {
       "type": "boolean",
       "default": true,
@@ -258,6 +271,11 @@
       "default": true,
       "description": "Enable or disable the plugin."
     },
+    "pylsp.plugins.rope_completion.eager": {
+      "type": "boolean",
+      "default": false,
+      "description": "Resolve documentation and detail eagerly."
+    },
     "pylsp.plugins.yapf.enabled": {
       "type": "boolean",
       "default": true,

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
@@ -38,7 +41,6 @@
 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)
     resolve_eagerly = settings.get('eager', False)
     code_position = _utils.position_to_jedi_linecolumn(document, position)
@@ -55,19 +57,34 @@ 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)
+    modules_to_cache_labels_for = settings.get('cache_labels_for', None)
+    if modules_to_cache_labels_for is not None:
+        LABEL_RESOLVER.cached_modules = modules_to_cache_labels_for
+
     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=resolve_eagerly,
+            resolve_label=(i < max_labels_resolve)
+        )
+        for i, c in enumerate(completions)
     ]
 
     # TODO split up once other improvements are merged
     if include_class_objects:
-        for c in completions:
+        for i, c in enumerate(completions):
             if c.type == 'class':
-                completion_dict = _format_completion(c, False, resolve=resolve_eagerly)
+                completion_dict = _format_completion(
+                    c,
+                    False,
+                    resolve=resolve_eagerly,
+                    resolve_label=(i < max_labels_resolve)
+                )
                 completion_dict['kind'] = lsp.CompletionItemKind.TypeParameter
                 completion_dict['label'] += ' object'
                 ready_completions.append(completion_dict)
@@ -150,9 +167,9 @@ def _resolve_completion(completion, d):
     return completion
 
 
-def _format_completion(d, include_params=True, resolve=False):
+def _format_completion(d, include_params=True, resolve=False, resolve_label=False):
     completion = {
-        'label': _label(d),
+        'label': _label(d, resolve_label),
         'kind': _TYPE_MAP.get(d.type),
         'sortText': _sort_text(d),
         'insertText': d.name
@@ -192,12 +209,12 @@ def _format_completion(d, include_params=True, resolve=False):
     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
 
 
@@ -216,3 +233,86 @@ 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'}
+
+    @property
+    def cached_modules(self):
+        return self._cached_modules
+
+    @cached_modules.setter
+    def cached_modules(self, new_value):
+        self._cached_modules = set(new_value)
+
+    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
 
@@ -148,6 +149,7 @@ def test_jedi_completion_item_resolve(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 for c in completions}
@@ -179,6 +181,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}
@@ -194,6 +214,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}