Create a cache for snippets (#83)

Author: ccordoba12

CONFIGURATION.md

@@ -22,8 +22,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_completion.resolve_at_most` | `number`  | How many labels and snippets (at most) should be resolved? | `25` |
+| `pylsp.plugins.jedi_completion.cache_for` | `array`  of  `string` items | Modules for which labels and snippets 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` |

pylsp/config/schema.json

@@ -104,18 +104,18 @@
       "default": false,
       "description": "Resolve documentation and detail eagerly."
     },
-    "pylsp.plugins.jedi_completion.resolve_at_most_labels": {
+    "pylsp.plugins.jedi_completion.resolve_at_most": {
       "type": "number",
       "default": 25,
-      "description": "How many labels (at most) should be resolved?"
+      "description": "How many labels and snippets (at most) should be resolved?"
     },
-    "pylsp.plugins.jedi_completion.cache_labels_for": {
+    "pylsp.plugins.jedi_completion.cache_for": {
       "type": "array",
       "items": {
         "type": "string"
       },
       "default": ["pandas", "numpy", "tensorflow", "matplotlib"],
-      "description": "Modules for which the labels should be cached."
+      "description": "Modules for which labels and snippets should be cached."
     },
     "pylsp.plugins.jedi_definition.enabled": {
       "type": "boolean",

pylsp/plugins/_resolvers.py

@@ -0,0 +1,135 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+from collections import defaultdict
+import logging
+from time import time
+
+from jedi.api.classes import Completion
+
+from pylsp import lsp
+
+
+log = logging.getLogger(__name__)
+
+
+# ---- Base class
+# -----------------------------------------------------------------------------
+class Resolver:
+
+    def __init__(self, callback, resolve_on_error, time_to_live=60 * 30):
+        self.callback = callback
+        self.resolve_on_error = resolve_on_error
+        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(completion)
+                self._cache_ttl[self.time_key()].add(key)
+            return self._cache[key]
+
+        return self.resolve(completion)
+
+    def _create_completion_id(self, completion: Completion):
+        return (
+            completion.full_name, completion.module_path,
+            completion.line, completion.column,
+            self.time_key()
+        )
+
+    def resolve(self, completion):
+        try:
+            sig = completion.get_signatures()
+            return self.callback(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 self.resolve_on_error
+
+
+# ---- Label resolver
+# -----------------------------------------------------------------------------
+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 = Resolver(callback=format_label, resolve_on_error='')
+
+
+# ---- Snippets resolver
+# -----------------------------------------------------------------------------
+def format_snippet(completion, sig):
+    if not sig:
+        return {}
+
+    snippet_completion = {}
+
+    positional_args = [param for param in sig[0].params
+                       if '=' not in param.description and
+                       param.name not in {'/', '*'}]
+
+    if len(positional_args) > 1:
+        # For completions with params, we can generate a snippet instead
+        snippet_completion['insertTextFormat'] = lsp.InsertTextFormat.Snippet
+        snippet = completion.name + '('
+        for i, param in enumerate(positional_args):
+            snippet += '${%s:%s}' % (i + 1, param.name)
+            if i < len(positional_args) - 1:
+                snippet += ', '
+        snippet += ')$0'
+        snippet_completion['insertText'] = snippet
+    elif len(positional_args) == 1:
+        snippet_completion['insertTextFormat'] = lsp.InsertTextFormat.Snippet
+        snippet_completion['insertText'] = completion.name + '($0)'
+    else:
+        snippet_completion['insertText'] = completion.name + '()'
+
+    return snippet_completion
+
+
+SNIPPET_RESOLVER = Resolver(callback=format_snippet, resolve_on_error={})

pylsp/plugins/jedi_completion.py

@@ -3,13 +3,11 @@
 
 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
+from pylsp.plugins._resolvers import LABEL_RESOLVER, SNIPPET_RESOLVER
 
 log = logging.getLogger(__name__)
 
@@ -57,10 +55,11 @@ 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
+    max_to_resolve = settings.get('resolve_at_most', 25)
+    modules_to_cache_for = settings.get('cache_for', None)
+    if modules_to_cache_for is not None:
+        LABEL_RESOLVER.cached_modules = modules_to_cache_for
+        SNIPPET_RESOLVER.cached_modules = modules_to_cache_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)
@@ -70,7 +69,7 @@ def pylsp_completions(config, document, position):
             c,
             include_params,
             resolve=resolve_eagerly,
-            resolve_label=(i < max_labels_resolve)
+            resolve_label_or_snippet=(i < max_to_resolve)
         )
         for i, c in enumerate(completions)
     ]
@@ -83,7 +82,7 @@ def pylsp_completions(config, document, position):
                     c,
                     False,
                     resolve=resolve_eagerly,
-                    resolve_label=(i < max_labels_resolve)
+                    resolve_label_or_snippet=(i < max_to_resolve)
                 )
                 completion_dict['kind'] = lsp.CompletionItemKind.TypeParameter
                 completion_dict['label'] += ' object'
@@ -175,9 +174,9 @@ def _resolve_completion(completion, d):
     return completion
 
 
-def _format_completion(d, include_params=True, resolve=False, resolve_label=False):
+def _format_completion(d, include_params=True, resolve=False, resolve_label_or_snippet=False):
     completion = {
-        'label': _label(d, resolve_label),
+        'label': _label(d, resolve_label_or_snippet),
         'kind': _TYPE_MAP.get(d.type),
         'sortText': _sort_text(d),
         'insertText': d.name
@@ -193,29 +192,8 @@ def _format_completion(d, include_params=True, resolve=False, resolve_label=Fals
         completion['insertText'] = path
 
     if include_params and not is_exception_class(d.name):
-        sig = d.get_signatures()
-        if not sig:
-            return completion
-
-        positional_args = [param for param in sig[0].params
-                           if '=' not in param.description and
-                           param.name not in {'/', '*'}]
-
-        if len(positional_args) > 1:
-            # For completions with params, we can generate a snippet instead
-            completion['insertTextFormat'] = lsp.InsertTextFormat.Snippet
-            snippet = d.name + '('
-            for i, param in enumerate(positional_args):
-                snippet += '${%s:%s}' % (i + 1, param.name)
-                if i < len(positional_args) - 1:
-                    snippet += ', '
-            snippet += ')$0'
-            completion['insertText'] = snippet
-        elif len(positional_args) == 1:
-            completion['insertTextFormat'] = lsp.InsertTextFormat.Snippet
-            completion['insertText'] = d.name + '($0)'
-        else:
-            completion['insertText'] = d.name + '()'
+        snippet = _snippet(d, resolve_label_or_snippet)
+        completion.update(snippet)
 
     return completion
 
@@ -229,6 +207,13 @@ def _label(definition, resolve=False):
     return definition.name
 
 
+def _snippet(definition, resolve=False):
+    if not resolve:
+        return {}
+    snippet = SNIPPET_RESOLVER.get_or_create(definition)
+    return snippet
+
+
 def _detail(definition):
     try:
         return definition.parent().full_name or ''
@@ -244,86 +229,3 @@ 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)