discovering and semantic parsing now runs in own threads

Author: d-biehl

package.json

@@ -137,7 +137,7 @@
         "configuration": "./language-configuration.json"
       }
     ],
-    "#grammars": [
+    "grammars": [
       {
         "language": "robotframework",
         "scopeName": "source.robotframework",

robotcode/language_server/robotframework/parts/discovering.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
+import ast
+import asyncio
+from itertools import chain
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterator, List, Optional
 
 from ....jsonrpc2.protocol import rpc_method
-from ....utils.async_itertools import async_chain
 from ....utils.logging import LoggingDescriptor
 from ....utils.uri import Uri
 from ...common.text_document import TextDocument
 from ...common.types import DocumentUri, Model, Position, Range, TextDocumentIdentifier
-from ...robotframework.utils.async_ast import walk
 from .protocol_part import RobotLanguageServerProtocolPart
 
 if TYPE_CHECKING:
@@ -62,8 +63,7 @@ def get_document(self, uri: DocumentUri) -> TextDocument:
 
         return TextDocument(document_uri=uri, language_id="robot", version=None, text=text)
 
-    @rpc_method(name="robot/discovering/getTestsFromWorkspace", param_type=GetAllTestsParams)
-    async def get_tests_from_workspace(self, paths: Optional[List[str]]) -> List[TestItem]:
+    def get_tests_from_workspace_threading(self, paths: Optional[List[str]]) -> List[TestItem]:
         from robot.output.logger import LOGGER
         from robot.running import TestCase, TestSuite
 
@@ -151,13 +151,16 @@ def nonexisting_paths(paths: List[str]) -> Iterator[str]:
             except BaseException as e:
                 return [TestItem(type="error", id=Path.cwd().name, label=Path.cwd().name, error=str(e))]
 
-    @rpc_method(name="robot/discovering/getTestsFromDocument", param_type=GetTestsParams)
-    async def get_tests_from_document(self, text_document: TextDocumentIdentifier, id: Optional[str]) -> List[TestItem]:
+    @rpc_method(name="robot/discovering/getTestsFromWorkspace", param_type=GetAllTestsParams)
+    async def get_tests_from_workspace(self, paths: Optional[List[str]]) -> List[TestItem]:
+        return await asyncio.get_event_loop().run_in_executor(None, self.get_tests_from_workspace_threading, paths)
+
+    def get_tests_from_document_threading(
+        self, text_document: TextDocumentIdentifier, id: Optional[str], model: ast.AST
+    ) -> List[TestItem]:
         from robot.parsing.model.blocks import TestCase
         from robot.parsing.model.statements import Tags
 
-        model = await self.parent.documents_cache.get_model(self.get_document(text_document.uri))
-
         return [
             TestItem(
                 type="test",
@@ -172,12 +175,19 @@ async def get_tests_from_document(self, text_document: TextDocumentIdentifier, i
                     ),
                 ),
                 tags=[
-                    str(tag)
-                    async for tag in async_chain(
-                        *[tags.values async for tags in walk(test_case) if isinstance(tags, Tags)]
-                    )
+                    str(tag) for tag in chain(*[tags.values for tags in ast.walk(test_case) if isinstance(tags, Tags)])
                 ],
             )
-            async for test_case in walk(model)
+            for test_case in ast.walk(model)
             if isinstance(test_case, TestCase)
         ]
+
+    @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 asyncio.get_event_loop().run_in_executor(
+            None,
+            self.get_tests_from_document_threading,
+            text_document,
+            id,
+            await self.parent.documents_cache.get_model(self.get_document(text_document.uri)),
+        )

robotcode/language_server/robotframework/parts/semantic_tokens.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
+import asyncio
 import operator
 import re
 from enum import Enum
 from functools import reduce
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncGenerator,
     Dict,
     FrozenSet,
+    Generator,
+    Iterable,
     NamedTuple,
     Optional,
     Set,
@@ -96,7 +98,7 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         parent.semantic_tokens.token_types += [e for e in RobotSemTokenTypes]
         parent.semantic_tokens.collect_full.add(self.collect_full)
         parent.semantic_tokens.collect_range.add(self.collect_range)
-        parent.semantic_tokens.collect_full_delta.add(self.collect_full_delta)
+        # parent.semantic_tokens.collect_full_delta.add(self.collect_full_delta)
 
     @classmethod
     def generate_mapping(cls) -> Dict[str, Tuple[Enum, Optional[Set[Enum]]]]:
@@ -164,9 +166,9 @@ def mapping(cls) -> Dict[str, Tuple[Enum, Optional[Set[Enum]]]]:
     )
 
     @classmethod
-    async def generate_sem_sub_tokens(
+    def generate_sem_sub_tokens(
         cls, token: Token, col_offset: Optional[int] = None, length: Optional[int] = None
-    ) -> AsyncGenerator[SemTokenInfo, None]:
+    ) -> Generator[SemTokenInfo, None, None]:
         from robot.parsing.lexer.tokens import Token as RobotToken
 
         sem_info = cls.mapping().get(token.type, None) if token.type is not None else None
@@ -199,55 +201,55 @@ async def generate_sem_sub_tokens(
                 yield SemTokenInfo.from_token(token, sem_info[0], sem_info[1], col_offset, length)
 
     @classmethod
-    async def generate_sem_tokens(cls, token: Token) -> AsyncGenerator[SemTokenInfo, None]:
+    def generate_sem_tokens(cls, token: Token) -> Generator[SemTokenInfo, None, None]:
         from robot.parsing.lexer.tokens import Token as RobotToken
 
         if token.type in RobotToken.ALLOW_VARIABLES:
             last_sub_token = token
             try:
                 for sub_token in token.tokenize_variables():
                     last_sub_token = sub_token
-                    async for e in cls.generate_sem_sub_tokens(sub_token):
+                    for e in cls.generate_sem_sub_tokens(sub_token):
                         yield e
             except BaseException:
                 pass
             if last_sub_token == token:
-                async for e in cls.generate_sem_sub_tokens(last_sub_token):
+                for e in cls.generate_sem_sub_tokens(last_sub_token):
                     yield e
             elif last_sub_token is not None and last_sub_token.end_col_offset < token.end_col_offset:
-                async for e in cls.generate_sem_sub_tokens(
+                for e in cls.generate_sem_sub_tokens(
                     token,
                     last_sub_token.end_col_offset,
                     token.end_col_offset - last_sub_token.end_col_offset - last_sub_token.col_offset,
                 ):
                     yield e
 
         else:
-            async for e in cls.generate_sem_sub_tokens(token):
+            for e in cls.generate_sem_sub_tokens(token):
                 yield e
 
-    async def collect(
-        self, document: TextDocument, range: Optional[Range]
+    def collect_threading(
+        self, tokens: Iterable[Token], range: Optional[Range]
     ) -> Union[SemanticTokens, SemanticTokensPartialResult, None]:
 
         data = []
         last_line = 0
         last_col = 0
 
-        tokens = await self.parent.documents_cache.get_tokens(document)
-
         start = True
+
         for robot_token in tokens:
             if range is not None:
-                if start and not token_in_range(robot_token, range):
-                    continue
+                if start:
+                    if not token_in_range(robot_token, range):
+                        continue
+                    else:
+                        start = False
                 else:
-                    start = False
-
-                if not start and not token_in_range(robot_token, range):
-                    break
+                    if not token_in_range(robot_token, range):
+                        break
 
-            async for token in self.generate_sem_tokens(robot_token):
+            for token in self.generate_sem_tokens(robot_token):
                 current_line = token.lineno - 1
 
                 data.append(current_line - last_line)
@@ -281,13 +283,18 @@ async def collect(
     async def collect_full(
         self, sender: Any, document: TextDocument, **kwargs: Any
     ) -> Union[SemanticTokens, SemanticTokensPartialResult, None]:
-        return await self.collect(document, None)
+
+        return await asyncio.get_event_loop().run_in_executor(
+            None, self.collect_threading, await self.parent.documents_cache.get_tokens(document), None
+        )
 
     @language_id("robotframework")
     async def collect_range(
         self, sender: Any, document: TextDocument, range: Range, **kwargs: Any
     ) -> Union[SemanticTokens, SemanticTokensPartialResult, None]:
-        return await self.collect(document, range)
+        return await asyncio.get_event_loop().run_in_executor(
+            None, self.collect_threading, await self.parent.documents_cache.get_tokens(document), range
+        )
 
     @language_id("robotframework")
     async def collect_full_delta(

robotcode/language_server/robotframework/utils/ast.py

@@ -43,9 +43,12 @@ def tokenize_variables(self) -> Iterator[Token]:
 
 
 @runtime_checkable
-class Statement(Protocol):
+class HasTokens(Protocol):
     tokens: Tuple[Token, ...]
 
+
+@runtime_checkable
+class Statement(Protocol):
     def get_token(self, type: str) -> Token:
         ...
 
@@ -142,7 +145,7 @@ def whitespace_from_begin_of_token(token: Token) -> str:
     return result
 
 
-def get_tokens_at_position(node: Statement, position: Position) -> List[Token]:
+def get_tokens_at_position(node: HasTokens, position: Position) -> List[Token]:
     return [t for t in node.tokens if position.is_in_range(range := range_from_token(t)) or range.end == position]