introduce cancelation_token to diagnostics, some renaming

d-biehl · d-biehl · commit 89e3eb64988d · 2021-11-29T13:02:10.000+01:00
diff --git a/robotcode/language_server/common/parts/diagnostics.py b/robotcode/language_server/common/parts/diagnostics.py
@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, cast
 
-from ....utils.async_tools import async_tasking_event_iterator
+from ....utils.async_tools import CancelationToken, async_tasking_event_iterator
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
 from ..language import language_id, language_id_filter
@@ -25,14 +25,21 @@
 class PublishDiagnosticsEntry:
     _logger = LoggingDescriptor()
 
-    def __init__(self, document_uri: DocumentUri, task_factory: Callable[..., asyncio.Task[Any]]) -> None:
+    def __init__(
+        self,
+        document_uri: DocumentUri,
+        cancelation_token: CancelationToken,
+        task_factory: Callable[..., asyncio.Task[Any]],
+    ) -> None:
 
         self._document_uri = document_uri
 
         self._task_factory = task_factory
 
         self._task: Optional[asyncio.Task[Any]] = None
 
+        self.cancel_token = cancelation_token
+
         @PublishDiagnosticsEntry._logger.call
         def create_task() -> None:
             self._task = self._task_factory()
@@ -74,6 +81,7 @@ async def cancel() -> None:
             t = self.task
             self._task = None
             if not t.done():
+                self.cancel_token.cancel()
                 t.cancel()
                 try:
                     await t
@@ -109,7 +117,9 @@ def __init__(self, protocol: LanguageServerProtocol) -> None:
         self.parent.documents.did_save.add(self.on_did_save)
 
     @async_tasking_event_iterator
-    async def collect(sender, document: TextDocument) -> DiagnosticsResult:  # NOSONAR
+    async def collect(
+        sender, document: TextDocument, cancelation_token: CancelationToken
+    ) -> DiagnosticsResult:  # NOSONAR
         ...
 
     @_logger.call
@@ -166,16 +176,17 @@ async def on_did_change(self, sender: Any, document: TextDocument) -> None:
     async def start_publish_diagnostics_task(self, document: TextDocument) -> None:
         async with self._task_lock:
             self._cancel_entry(self._running_diagnostics.get(document.uri, None))
-
+            cancelation_token = CancelationToken()
             self._running_diagnostics[document.uri] = PublishDiagnosticsEntry(
                 document.document_uri,
+                cancelation_token,
                 lambda: asyncio.create_task(
-                    self.publish_diagnostics(document.document_uri),
+                    self.publish_diagnostics(document.document_uri, cancelation_token),
                 ),
             )
 
     @_logger.call
-    async def publish_diagnostics(self, document_uri: DocumentUri) -> None:
+    async def publish_diagnostics(self, document_uri: DocumentUri, cancelation_token: CancelationToken) -> None:
         document = self.parent.documents.get(document_uri, None)
         if document is None:
             return
@@ -187,6 +198,7 @@ async def publish_diagnostics(self, document_uri: DocumentUri) -> None:
         async for result_any in self.collect(
             self,
             document,
+            cancelation_token,
             callback_filter=language_id_filter(document),
             return_exceptions=True,
         ):
diff --git a/robotcode/language_server/robotframework/diagnostics/analyzer.py b/robotcode/language_server/robotframework/diagnostics/analyzer.py
@@ -4,6 +4,7 @@
 import asyncio
 from typing import List, Optional, cast
 
+from ....utils.async_tools import CancelationToken
 from ....utils.uri import Uri
 from ...common.lsp_types import (
     Diagnostic,
@@ -26,16 +27,24 @@
 
 
 class Analyzer(AsyncVisitor):
-    async def get(self, model: ast.AST, namespace: Namespace) -> List[Diagnostic]:
+    async def get(
+        self, model: ast.AST, namespace: Namespace, cancelation_token: Optional[CancelationToken] = None
+    ) -> List[Diagnostic]:
         self._results: List[Diagnostic] = []
         self._namespace = namespace
-
+        self.cancelation_token = cancelation_token
         self.current_testcase_or_keyword_name: Optional[str] = None
         self.finder = KeywordFinder(self._namespace)
 
         await self.visit(model)
         return self._results
 
+    async def visit(self, node: ast.AST) -> None:
+        if self.cancelation_token:
+            self.cancelation_token.throw_if_canceled()
+
+        await super().visit(node)
+
     async def _analyze_keyword_call(
         self,
         keyword: Optional[str],
diff --git a/robotcode/language_server/robotframework/diagnostics/namespace.py b/robotcode/language_server/robotframework/diagnostics/namespace.py
@@ -23,7 +23,7 @@
 )
 
 from ....utils.async_itertools import async_chain
-from ....utils.async_tools import awaitable_to_thread
+from ....utils.async_tools import CancelationToken, awaitable_run_in_thread
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
 from ...common.lsp_types import (
@@ -577,10 +577,10 @@ async def invalidate(self) -> None:
         self.invalidated_callback(self)
 
     @_logger.call
-    async def get_diagnostisc(self) -> List[Diagnostic]:
+    async def get_diagnostisc(self, cancelation_token: Optional[CancelationToken] = None) -> List[Diagnostic]:
         await self.ensure_initialized()
 
-        await self._analyze()
+        await self._analyze(cancelation_token)
 
         return self._diagnostics
 
@@ -1088,14 +1088,16 @@ async def get_keywords(self) -> List[KeywordDoc]:
         return self._keywords
 
     @_logger.call
-    async def _analyze(self) -> None:
+    async def _analyze(self, cancelation_token: Optional[CancelationToken] = None) -> None:
         from .analyzer import Analyzer
 
         if not self._analyzed:
             async with self._analyze_lock:
                 if not self._analyzed:
                     try:
-                        self._diagnostics += await awaitable_to_thread(Analyzer().get(self.model, self))
+                        self._diagnostics += await awaitable_run_in_thread(
+                            Analyzer().get(self.model, self, cancelation_token)
+                        )
 
                         lib_doc = await self.get_library_doc()
 
diff --git a/robotcode/language_server/robotframework/parts/diagnostics.py b/robotcode/language_server/robotframework/parts/diagnostics.py
@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 import ast
+import asyncio
 from typing import TYPE_CHECKING, Any, List, Optional
 
-from ....utils.async_tools import awaitable_to_thread
+from ....utils.async_tools import CancelationToken, awaitable_run_in_thread
 from ....utils.logging import LoggingDescriptor
 from ...common.language import language_id
 from ...common.lsp_types import Diagnostic, DiagnosticSeverity, Position, Range
@@ -60,14 +61,18 @@ def _create_error_from_token(self, token: Token, source: Optional[str] = None) -
 
     @language_id("robotframework")
     @_logger.call
-    async def collect_token_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_token_errors(
+        self, sender: Any, document: TextDocument, cancelation_token: CancelationToken
+    ) -> DiagnosticsResult:
         async def collect_async(tokens: List[Token]) -> List[Diagnostic]:
             from robot.errors import VariableError
             from robot.parsing.lexer.tokens import Token
 
             result: List[Diagnostic] = []
             try:
                 for token in tokens:
+                    cancelation_token.throw_if_canceled()
+
                     if token.type in [Token.ERROR, Token.FATAL_ERROR]:
                         result.append(self._create_error_from_token(token))
 
@@ -89,8 +94,7 @@ async def collect_async(tokens: List[Token]) -> List[Diagnostic]:
                                 code=type(e).__qualname__,
                             )
                         )
-
-            except (SystemExit, KeyboardInterrupt):
+            except (asyncio.CancelledError, SystemExit, KeyboardInterrupt):
                 raise
             except BaseException as e:
                 return [
@@ -116,18 +120,22 @@ async def collect_async(tokens: List[Token]) -> List[Diagnostic]:
 
         return DiagnosticsResult(
             self.collect_token_errors,
-            await awaitable_to_thread(collect_async(await self.parent.documents_cache.get_tokens(document))),
+            await awaitable_run_in_thread(collect_async(await self.parent.documents_cache.get_tokens(document))),
         )
 
     @language_id("robotframework")
     @_logger.call
-    async def collect_walk_model_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_walk_model_errors(
+        self, sender: Any, document: TextDocument, cancelation_token: CancelationToken
+    ) -> DiagnosticsResult:
         async def collect_async(model: ast.AST) -> List[Diagnostic]:
             from ..utils.ast import HasError, HasErrors
             from ..utils.async_ast import iter_nodes
 
             result: List[Diagnostic] = []
             async for node in iter_nodes(model):
+                cancelation_token.throw_if_canceled()
+
                 error = node.error if isinstance(node, HasError) else None
                 if error is not None:
                     result.append(self._create_error_from_node(node, error))
@@ -139,14 +147,16 @@ async def collect_async(model: ast.AST) -> List[Diagnostic]:
 
         return DiagnosticsResult(
             self.collect_walk_model_errors,
-            await awaitable_to_thread(collect_async(await self.parent.documents_cache.get_model(document))),
+            await awaitable_run_in_thread(collect_async(await self.parent.documents_cache.get_model(document))),
         )
 
     @language_id("robotframework")
     @_logger.call
-    async def collect_namespace_diagnostics(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_namespace_diagnostics(
+        self, sender: Any, document: TextDocument, cancelation_token: CancelationToken
+    ) -> DiagnosticsResult:
         namespace = await self.parent.documents_cache.get_namespace(document)
         if namespace is None:
             return DiagnosticsResult(self.collect_namespace_diagnostics, None)
 
-        return DiagnosticsResult(self.collect_namespace_diagnostics, await namespace.get_diagnostisc())
+        return DiagnosticsResult(self.collect_namespace_diagnostics, await namespace.get_diagnostisc(cancelation_token))
diff --git a/robotcode/language_server/robotframework/parts/discovering.py b/robotcode/language_server/robotframework/parts/discovering.py
@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING, Iterator, List, Optional
 
 from ....jsonrpc2.protocol import rpc_method
-from ....utils.async_tools import to_thread
+from ....utils.async_tools import run_in_thread
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
 from ...common.lsp_types import (
@@ -163,7 +163,7 @@ def nonexisting_paths(paths: List[str]) -> Iterator[str]:
 
     @rpc_method(name="robot/discovering/getTestsFromWorkspace", param_type=GetAllTestsParams)
     async def get_tests_from_workspace(self, paths: Optional[List[str]]) -> List[TestItem]:
-        return await to_thread(self.get_tests_from_workspace_threading, paths)
+        return await run_in_thread(self.get_tests_from_workspace_threading, paths)
 
     def get_tests_from_document_threading(
         self, text_document: TextDocumentIdentifier, id: Optional[str], model: ast.AST
@@ -194,7 +194,7 @@ def get_tests_from_document_threading(
 
     @rpc_method(name="robot/discovering/getTestsFromDocument", param_type=GetTestsParams)
     async def get_tests_from_document(self, text_document: TextDocumentIdentifier, id: Optional[str]) -> List[TestItem]:
-        return await to_thread(
+        return await run_in_thread(
             self.get_tests_from_document_threading,
             text_document,
             id,
diff --git a/robotcode/language_server/robotframework/parts/documents_cache.py b/robotcode/language_server/robotframework/parts/documents_cache.py
@@ -16,7 +16,7 @@
     cast,
 )
 
-from ....utils.async_tools import CancelationToken, async_tasking_event, to_thread
+from ....utils.async_tools import CancelationToken, async_tasking_event, run_in_thread
 from ....utils.uri import Uri
 from ...common.language import language_id_filter
 from ...common.parts.workspace import WorkspaceFolder
@@ -111,7 +111,7 @@ async def __get_tokens_internal(
         try:
             if cancelation_token is None:
                 cancelation_token = CancelationToken()
-            return await to_thread(get, document.text, cancelation_token)
+            return await run_in_thread(get, document.text, cancelation_token)
         except asyncio.CancelledError:
             if cancelation_token is not None:
                 cancelation_token.cancel()
@@ -181,7 +181,7 @@ def get_tokens(_source: str, _data_only: bool = False) -> Generator[Token, None,
                 yield t
 
         try:
-            model = await to_thread(_get_model, get_tokens, document.uri.to_path())
+            model = await run_in_thread(_get_model, get_tokens, document.uri.to_path())
         except asyncio.CancelledError:
             if cancelation_token is not None:
                 cancelation_token.cancel()
diff --git a/robotcode/language_server/robotframework/parts/references.py b/robotcode/language_server/robotframework/parts/references.py
@@ -16,7 +16,7 @@
     cast,
 )
 
-from ....utils.async_tools import CancelationToken, awaitable_to_thread
+from ....utils.async_tools import CancelationToken, awaitable_run_in_thread
 from ....utils.glob_path import iter_files
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
@@ -305,7 +305,9 @@ async def _find_keyword_references_in_file(
         ):
             return []
 
-        return await awaitable_to_thread(self._find_keyword_references_in_namespace(namespace, kw_doc, cancel_token))
+        return await awaitable_run_in_thread(
+            self._find_keyword_references_in_namespace(namespace, kw_doc, cancel_token)
+        )
 
     async def _find_keyword_references_in_namespace(
         self, namespace: Namespace, kw_doc: KeywordDoc, cancel_token: CancelationToken
diff --git a/robotcode/language_server/robotframework/parts/robocop_diagnostics.py b/robotcode/language_server/robotframework/parts/robocop_diagnostics.py
@@ -5,7 +5,7 @@
 import io
 from typing import TYPE_CHECKING, Any, List, Optional
 
-from ....utils.async_tools import to_thread
+from ....utils.async_tools import CancelationToken, run_in_thread
 from ....utils.logging import LoggingDescriptor
 from ...common.language import language_id
 from ...common.lsp_types import Diagnostic, DiagnosticSeverity, Position, Range
@@ -48,7 +48,9 @@ async def get_config(self, document: TextDocument) -> Optional[RoboCopConfig]:
 
     @language_id("robotframework")
     @_logger.call
-    async def collect_diagnostics(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_diagnostics(
+        self, sender: Any, document: TextDocument, cancelation_token: CancelationToken
+    ) -> DiagnosticsResult:
 
         try:
             workspace_folder = self.parent.workspace.get_workspace_folder(document.uri)
@@ -58,7 +60,9 @@ async def collect_diagnostics(self, sender: Any, document: TextDocument) -> Diag
                 if extension_config is not None and extension_config.enabled:
 
                     model = await self.parent.documents_cache.get_model(document)
-                    result = await self.collect_threading(document, workspace_folder, extension_config, model)
+                    result = await self.collect_threading(
+                        document, workspace_folder, extension_config, model, cancelation_token
+                    )
                     return DiagnosticsResult(self.collect_diagnostics, result)
         except (SystemExit, KeyboardInterrupt, asyncio.CancelledError):
             raise
@@ -68,12 +72,22 @@ async def collect_diagnostics(self, sender: Any, document: TextDocument) -> Diag
         return DiagnosticsResult(self.collect_diagnostics, [])
 
     async def collect_threading(
-        self, document: TextDocument, workspace_folder: WorkspaceFolder, extension_config: RoboCopConfig, model: ast.AST
+        self,
+        document: TextDocument,
+        workspace_folder: WorkspaceFolder,
+        extension_config: RoboCopConfig,
+        model: ast.AST,
+        cancelation_token: CancelationToken,
     ) -> List[Diagnostic]:
-        return await to_thread(self.collect, document, workspace_folder, extension_config, model)
+        return await run_in_thread(self.collect, document, workspace_folder, extension_config, model, cancelation_token)
 
     def collect(
-        self, document: TextDocument, workspace_folder: WorkspaceFolder, extension_config: RoboCopConfig, model: ast.AST
+        self,
+        document: TextDocument,
+        workspace_folder: WorkspaceFolder,
+        extension_config: RoboCopConfig,
+        model: ast.AST,
+        cancelation_token: CancelationToken,
     ) -> List[Diagnostic]:
         from robocop.config import Config
         from robocop.rules import RuleSeverity
@@ -98,6 +112,7 @@ def collect(
             analyser = Robocop(from_cli=False, config=config)
             analyser.reload_config()
 
+            # TODO find a way to cancel the run_check
             issues = analyser.run_check(model, str(document.uri.to_path()), document.text)
 
             for issue in issues:
@@ -117,7 +132,7 @@ def collect(
                     source=self.source_name,
                     code=f"{issue.severity.value}{issue.rule_id}",
                 )
-
+                cancelation_token.throw_if_canceled()
                 result.append(d)
 
         return result
diff --git a/robotcode/language_server/robotframework/parts/semantic_tokens.py b/robotcode/language_server/robotframework/parts/semantic_tokens.py
diff --git a/robotcode/utils/async_tools.py b/robotcode/utils/async_tools.py
