Report autoimport progress

pylsp/plugins/_rope_task_handle.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import logging
+from typing import Callable, ContextManager, List, Optional, Sequence
+
+from rope.base.taskhandle import BaseJobSet, BaseTaskHandle
+
+from pylsp.workspace import Workspace
+
+log = logging.getLogger(__name__)
+Report = Callable[[str, int], None]
+
+
+class PylspJobSet(BaseJobSet):
+    count: int = 0
+    done: int = 0
+    _reporter: Report
+    _report_iter: ContextManager
+    job_name: str = ""
+
+    def __init__(self, count: Optional[int], report_iter: ContextManager):
+        if count is not None:
+            self.count = count
+        self._reporter = report_iter.__enter__()
+        self._report_iter = report_iter
+
+    def started_job(self, name: Optional[str]) -> None:
+        if name:
+            self.job_name = name
+
+    def finished_job(self) -> None:
+        self.done += 1
+        if self.get_percent_done() is not None and int(self.get_percent_done()) >= 100:
+            if self._report_iter is None:
+                return
+            self._report_iter.__exit__(None, None, None)
+            self._report_iter = None
+        else:
+            self._report()
+
+    def check_status(self) -> None:
+        pass
+
+    def get_percent_done(self) -> Optional[float]:
+        if self.count == 0:
+            return 0
+        return (self.done / self.count) * 100
+
+    def increment(self) -> None:
+        """
+        Increment the number of tasks to complete.
+
+        This is used if the number is not known ahead of time.
+        """
+        self.count += 1
+        self._report()
+
+    def _report(self):
+        percent = int(self.get_percent_done())
+        message = f"{self.job_name} {self.done}/{self.count}"
+        log.debug(f"Reporting {message} {percent}%")
+        self._reporter(message, percent)
+
+
+class PylspTaskHandle(BaseTaskHandle):
+    name: str
+    observers: List
+    job_sets: List[PylspJobSet]
+    stopped: bool
+    workspace: Workspace
+    _report: Callable[[str, str], None]
+
+    def __init__(self, workspace: Workspace):
+        self.workspace = workspace
+        self.job_sets = []
+        self.observers = []
+
+    def create_jobset(self, name="JobSet", count: Optional[int] = None):
+        report_iter = self.workspace.report_progress(name, None, None)
+        result = PylspJobSet(count, report_iter)
+        self.job_sets.append(result)
+        self._inform_observers()
+        return result
+
+    def stop(self) -> None:
+        pass
+
+    def current_jobset(self) -> Optional[BaseJobSet]:
+        pass
+
+    def add_observer(self) -> None:
+        pass
+
+    def is_stopped(self) -> bool:
+        pass
+
+    def get_jobsets(self) -> Sequence[BaseJobSet]:
+        pass
+
+    def _inform_observers(self) -> None:
+        for observer in self.observers:
+            observer()

pylsp/plugins/rope_autoimport.py

@@ -1,7 +1,7 @@
 # Copyright 2022- Python Language Server Contributors.
 
 import logging
-from typing import Any, Dict, Generator, List, Set
+from typing import Any, Dict, Generator, List, Optional, Set
 
 import parso
 from jedi import Script
@@ -15,6 +15,8 @@
 from pylsp.config.config import Config
 from pylsp.workspace import Document, Workspace
 
+from ._rope_task_handle import PylspTaskHandle
+
 log = logging.getLogger(__name__)
 
 _score_pow = 5
@@ -44,8 +46,10 @@ def _should_insert(expr: tree.BaseNode, word_node: tree.Leaf) -> bool:
     if first_child == word_node:
         return True  # If the word is the first word then its fine
     if len(expr.children) > 1:
-        if any(node.type == "operator" and "." in node.value or
-               node.type == "trailer" for node in expr.children):
+        if any(
+            node.type == "operator" and "." in node.value or node.type == "trailer"
+            for node in expr.children
+        ):
             return False  # Check if we're on a method of a function
     if isinstance(first_child, (tree.PythonErrorNode, tree.PythonNode)):
         # The tree will often include error nodes like this to indicate errors
@@ -54,8 +58,7 @@ def _should_insert(expr: tree.BaseNode, word_node: tree.Leaf) -> bool:
     return _handle_first_child(first_child, expr, word_node)
 
 
-def _handle_first_child(first_child: NodeOrLeaf, expr: tree.BaseNode,
-                        word_node: tree.Leaf) -> bool:
+def _handle_first_child(first_child: NodeOrLeaf, expr: tree.BaseNode, word_node: tree.Leaf) -> bool:
     """Check if we suggest imports given the following first child."""
     if isinstance(first_child, tree.Import):
         return False
@@ -119,22 +122,16 @@ def _process_statements(
         insert_line = autoimport.find_insertion_line(document.source) - 1
         start = {"line": insert_line, "character": 0}
         edit_range = {"start": start, "end": start}
-        edit = {
-            "range": edit_range,
-            "newText": suggestion.import_statement + "\n"
-        }
-        score = _get_score(suggestion.source, suggestion.import_statement,
-                           suggestion.name, word)
+        edit = {"range": edit_range, "newText": suggestion.import_statement + "\n"}
+        score = _get_score(suggestion.source, suggestion.import_statement, suggestion.name, word)
         if score > _score_max:
             continue
         # TODO make this markdown
         yield {
             "label": suggestion.name,
             "kind": suggestion.itemkind,
             "sortText": _sort_import(score),
-            "data": {
-                "doc_uri": doc_uri
-            },
+            "data": {"doc_uri": doc_uri},
             "detail": _document(suggestion.import_statement),
             "additionalTextEdits": [edit],
         }
@@ -148,8 +145,7 @@ def get_names(script: Script) -> Set[str]:
 
 
 @hookimpl
-def pylsp_completions(config: Config, workspace: Workspace, document: Document,
-                      position):
+def pylsp_completions(config: Config, workspace: Workspace, document: Document, position):
     """Get autoimport suggestions."""
     line = document.lines[position["line"]]
     expr = parso.parse(line)
@@ -159,17 +155,15 @@ def pylsp_completions(config: Config, workspace: Workspace, document: Document,
     word = word_node.value
     log.debug(f"autoimport: searching for word: {word}")
     rope_config = config.settings(document_path=document.path).get("rope", {})
-    ignored_names: Set[str] = get_names(
-        document.jedi_script(use_document_path=True))
+    ignored_names: Set[str] = get_names(document.jedi_script(use_document_path=True))
     autoimport = workspace._rope_autoimport(rope_config)
-    suggestions = list(
-        autoimport.search_full(word, ignored_names=ignored_names))
+    suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
     results = list(
         sorted(
-            _process_statements(suggestions, document.uri, word, autoimport,
-                                document),
+            _process_statements(suggestions, document.uri, word, autoimport, document),
             key=lambda statement: statement["sortText"],
-        ))
+        )
+    )
     if len(results) > MAX_RESULTS:
         results = results[:MAX_RESULTS]
     return results
@@ -179,11 +173,10 @@ def _document(import_statement: str) -> str:
     return """# Auto-Import\n""" + import_statement
 
 
-def _get_score(source: int, full_statement: str, suggested_name: str,
-               desired_name) -> int:
+def _get_score(source: int, full_statement: str, suggested_name: str, desired_name) -> int:
     import_length = len("import")
     full_statement_score = len(full_statement) - import_length
-    suggested_name_score = ((len(suggested_name) - len(desired_name)))**2
+    suggested_name_score = ((len(suggested_name) - len(desired_name))) ** 2
     source_score = 20 * source
     return suggested_name_score + full_statement_score + source_score
 
@@ -196,24 +189,37 @@ def _sort_import(score: int) -> str:
     return "[z" + str(score).rjust(_score_pow, "0")
 
 
-@hookimpl
-def pylsp_initialize(config: Config, workspace: Workspace):
-    """Initialize AutoImport. Generates the cache for local and global items."""
-    memory: bool = config.plugin_settings("rope_autoimport").get(
-        "memory", False)
+def _reload_cache(config: Config, workspace: Workspace, files: Optional[List[Document]] = None):
+    memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
     rope_config = config.settings().get("rope", {})
     autoimport = workspace._rope_autoimport(rope_config, memory)
-    autoimport.generate_modules_cache()
-    autoimport.generate_cache()
+    task_handle = PylspTaskHandle(workspace)
+    resources: Optional[List[Resource]] = (
+        None if files is None else [document._rope_resource(rope_config) for document in files]
+    )
+    autoimport.generate_cache(task_handle=task_handle, resources=resources)
+    autoimport.generate_modules_cache(task_handle=task_handle)
 
 
 @hookimpl
-def pylsp_document_did_save(config: Config, workspace: Workspace,
-                            document: Document):
+def pylsp_initialize(config: Config, workspace: Workspace):
+    """Initialize AutoImport.
+
+    Generates the cache for local and global items.
+    """
+    _reload_cache(config, workspace)
+
+
+@hookimpl
+def pylsp_document_did_open(config: Config, workspace: Workspace):
+    """Initialize AutoImport.
+
+    Generates the cache for local and global items.
+    """
+    _reload_cache(config, workspace)
+
+
+@hookimpl
+def pylsp_document_did_save(config: Config, workspace: Workspace, document: Document):
     """Update the names associated with this document."""
-    rope_config = config.settings().get("rope", {})
-    rope_doucment: Resource = document._rope_resource(rope_config)
-    autoimport = workspace._rope_autoimport(rope_config)
-    autoimport.generate_cache(resources=[rope_doucment])
-    # Might as well using saving the document as an indicator to regenerate the module cache
-    autoimport.generate_modules_cache()
+    _reload_cache(config, workspace, [document])