Optimize collecting model errors, fixes #42

d-biehl · d-biehl · commit 57551f50f99d · 2022-03-13T11:42:39.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,10 @@ All notable changes to the "robotcode" extension will be documented in this file
 
 ## [Unreleased]
 
-- none so far
+### added
+
+- Optimize collecting model errors
+  - also fixes [#42](https://github.com/d-biehl/robotcode/issues/42)
 
 ##  0.8.0
 
diff --git a/robotcode/language_server/robotframework/parts/completion.py b/robotcode/language_server/robotframework/parts/completion.py
@@ -108,8 +108,10 @@ async def collect(
         if namespace is None:
             return None
 
+        model = await self.parent.documents_cache.get_model(document, False)
+
         return await CompletionCollector(
-            self.parent, document, namespace, await self.get_section_style(document)
+            self.parent, document, model, namespace, await self.get_section_style(document)
         ).collect(
             position,
             context,
@@ -126,9 +128,10 @@ async def resolve(self, sender: Any, completion_item: CompletionItem) -> Complet
                 document = await self.parent.documents.get(document_uri)
                 if document is not None:
                     namespace = await self.parent.documents_cache.get_namespace(document)
+                    model = await self.parent.documents_cache.get_model(document, False)
                     if namespace is not None:
                         return await CompletionCollector(
-                            self.parent, document, namespace, await self.get_section_style(document)
+                            self.parent, document, model, namespace, await self.get_section_style(document)
                         ).resolve(completion_item)
 
         return completion_item
@@ -221,11 +224,17 @@ class CompletionCollector(ModelHelperMixin):
     _logger = LoggingDescriptor()
 
     def __init__(
-        self, parent: RobotLanguageServerProtocol, document: TextDocument, namespace: Namespace, section_style: str
+        self,
+        parent: RobotLanguageServerProtocol,
+        document: TextDocument,
+        model: ast.AST,
+        namespace: Namespace,
+        section_style: str,
     ) -> None:
         self.parent = parent
         self.section_style = section_style
         self.document = document
+        self.model = model
         self.namespace = namespace
 
     async def _find_methods(self, cls: Type[Any]) -> AsyncGenerator[_CompleteMethod, None]:
@@ -248,7 +257,7 @@ async def collect(
         self, position: Position, context: Optional[CompletionContext]
     ) -> Union[List[CompletionItem], CompletionList, None]:
 
-        result_nodes = await get_nodes_at_position(self.namespace.model, position)
+        result_nodes = await get_nodes_at_position(self.model, position)
 
         result_nodes.reverse()
 
diff --git a/robotcode/language_server/robotframework/parts/diagnostics.py b/robotcode/language_server/robotframework/parts/diagnostics.py
@@ -11,7 +11,7 @@
 from ...common.parts.diagnostics import DiagnosticsResult
 from ...common.text_document import TextDocument
 from ..diagnostics.analyzer import Analyzer
-from ..utils.ast import Token, range_from_node, range_from_token
+from ..utils.ast import HeaderAndBodyBlock, Token, range_from_node, range_from_token
 
 if TYPE_CHECKING:
     from ..protocol import RobotLanguageServerProtocol
@@ -28,7 +28,7 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         self.source_name = "robotcode.diagnostics"
 
         parent.diagnostics.collect.add(self.collect_token_errors)
-        parent.diagnostics.collect.add(self.collect_walk_model_errors)
+        parent.diagnostics.collect.add(self.collect_model_errors)
 
         parent.diagnostics.collect.add(self.collect_namespace_diagnostics)
 
@@ -72,9 +72,21 @@ async def collect_namespace_diagnostics(self, sender: Any, document: TextDocumen
                 ],
             )
 
-    def _create_error_from_node(self, node: ast.AST, msg: str, source: Optional[str] = None) -> Diagnostic:
+    def _create_error_from_node(
+        self, node: ast.AST, msg: str, source: Optional[str] = None, only_start: bool = True
+    ) -> Diagnostic:
+        from robot.parsing.model.statements import Statement
+
+        if isinstance(node, HeaderAndBodyBlock):
+            if node.header is not None:
+                node = node.header
+            elif node.body:
+                stmt = next((n for n in node.body if isinstance(n, Statement)), None)
+                if stmt is not None:
+                    node = stmt
+
         return Diagnostic(
-            range=range_from_node(node, True),
+            range=range_from_node(node, True, only_start),
             message=msg,
             severity=DiagnosticSeverity.ERROR,
             source=source if source is not None else self.source_name,
@@ -158,13 +170,13 @@ async def collect_token_errors(self, sender: Any, document: TextDocument) -> Dia
     @language_id("robotframework")
     @threaded()
     @_logger.call
-    async def collect_walk_model_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
+    async def collect_model_errors(self, sender: Any, document: TextDocument) -> DiagnosticsResult:
 
         from ..utils.ast import HasError, HasErrors
         from ..utils.async_ast import iter_nodes
 
         try:
-            model = await self.parent.documents_cache.get_model(document)
+            model = await self.parent.documents_cache.get_model(document, True)
 
             result: List[Diagnostic] = []
             async for node in iter_nodes(model):
@@ -177,13 +189,13 @@ async def collect_walk_model_errors(self, sender: Any, document: TextDocument) -
                         if not await Analyzer.should_ignore(document, range_from_node(node)):
                             result.append(self._create_error_from_node(node, e))
 
-            return DiagnosticsResult(self.collect_walk_model_errors, result)
+            return DiagnosticsResult(self.collect_model_errors, result)
 
         except (asyncio.CancelledError, SystemExit, KeyboardInterrupt):
             raise
         except BaseException as e:
             return DiagnosticsResult(
-                self.collect_walk_model_errors,
+                self.collect_model_errors,
                 [
                     Diagnostic(
                         range=Range(
diff --git a/robotcode/language_server/robotframework/parts/documents_cache.py b/robotcode/language_server/robotframework/parts/documents_cache.py
@@ -63,9 +63,22 @@ async def __get_document_type(self, document: TextDocument) -> DocumentType:
         else:
             return DocumentType.UNKNOWN
 
-    async def get_tokens(self, document: TextDocument) -> List[Token]:
+    async def get_tokens(self, document: TextDocument, data_only: bool = False) -> List[Token]:
+        if data_only:
+            return await document.get_cache(self.__get_tokens_data_only)
         return await document.get_cache(self.__get_tokens)
 
+    async def __get_tokens_data_only(self, document: TextDocument) -> List[Token]:
+        document_type = await self.get_document_type(document)
+        if document_type == DocumentType.INIT:
+            return await self.get_init_tokens(document, True)
+        elif document_type == DocumentType.GENERAL:
+            return await self.get_general_tokens(document, True)
+        elif document_type == DocumentType.RESOURCE:
+            return await self.get_resource_tokens(document, True)
+        else:
+            raise UnknownFileTypeError(str(document.uri))
+
     async def __get_tokens(self, document: TextDocument) -> List[Token]:
         document_type = await self.get_document_type(document)
         if document_type == DocumentType.INIT:
@@ -77,9 +90,20 @@ async def __get_tokens(self, document: TextDocument) -> List[Token]:
         else:
             raise UnknownFileTypeError(str(document.uri))
 
-    async def get_general_tokens(self, document: TextDocument) -> List[Token]:
+    async def get_general_tokens(self, document: TextDocument, data_only: bool = False) -> List[Token]:
+        if data_only:
+            return await document.get_cache(self.__get_general_tokens_data_only)
         return await document.get_cache(self.__get_general_tokens)
 
+    async def __get_general_tokens_data_only(self, document: TextDocument) -> List[Token]:
+        import robot.api
+
+        def get(text: str) -> List[Token]:
+            with io.StringIO(text) as content:
+                return [e for e in robot.api.get_tokens(content, True) if check_canceled_sync()]
+
+        return await self.__get_tokens_internal(document, get)
+
     async def __get_general_tokens(self, document: TextDocument) -> List[Token]:
         import robot.api
 
@@ -97,9 +121,21 @@ async def __get_tokens_internal(
 
         return get(await document.text())
 
-    async def get_resource_tokens(self, document: TextDocument) -> List[Token]:
+    async def get_resource_tokens(self, document: TextDocument, data_only: bool = False) -> List[Token]:
+        if data_only:
+            return await document.get_cache(self.__get_resource_tokens_data_only)
+
         return await document.get_cache(self.__get_resource_tokens)
 
+    async def __get_resource_tokens_data_only(self, document: TextDocument) -> List[Token]:
+        import robot.api
+
+        def get(text: str) -> List[Token]:
+            with io.StringIO(text) as content:
+                return [e for e in robot.api.get_resource_tokens(content, True) if check_canceled_sync()]
+
+        return await self.__get_tokens_internal(document, get)
+
     async def __get_resource_tokens(self, document: TextDocument) -> List[Token]:
         import robot.api
 
@@ -109,9 +145,20 @@ def get(text: str) -> List[Token]:
 
         return await self.__get_tokens_internal(document, get)
 
-    async def get_init_tokens(self, document: TextDocument) -> List[Token]:
+    async def get_init_tokens(self, document: TextDocument, data_only: bool = False) -> List[Token]:
+        if data_only:
+            return await document.get_cache(self.__get_init_tokens_data_only)
         return await document.get_cache(self.__get_init_tokens)
 
+    async def __get_init_tokens_data_only(self, document: TextDocument) -> List[Token]:
+        import robot.api
+
+        def get(text: str) -> List[Token]:
+            with io.StringIO(text) as content:
+                return [e for e in robot.api.get_init_tokens(content, True) if check_canceled_sync()]
+
+        return await self.__get_tokens_internal(document, get)
+
     async def __get_init_tokens(self, document: TextDocument) -> List[Token]:
         import robot.api
 
@@ -121,15 +168,15 @@ def get(text: str) -> List[Token]:
 
         return await self.__get_tokens_internal(document, get)
 
-    async def get_model(self, document: TextDocument) -> ast.AST:
+    async def get_model(self, document: TextDocument, data_only: bool = True) -> ast.AST:
         document_type = await self.get_document_type(document)
 
         if document_type == DocumentType.INIT:
-            return await self.get_init_model(document)
+            return await self.get_init_model(document, data_only)
         if document_type == DocumentType.GENERAL:
-            return await self.get_general_model(document)
+            return await self.get_general_model(document, data_only)
         if document_type == DocumentType.RESOURCE:
-            return await self.get_resource_model(document)
+            return await self.get_resource_model(document, data_only)
         else:
             raise UnknownFileTypeError(f"Unknown file type '{document.uri}'.")
 
@@ -150,21 +197,37 @@ def get_tokens(_source: str, _data_only: bool = False) -> Generator[Token, None,
 
         return cast(ast.AST, model)
 
-    async def get_general_model(self, document: TextDocument) -> ast.AST:
+    async def get_general_model(self, document: TextDocument, data_only: bool = True) -> ast.AST:
+        if data_only:
+            return await document.get_cache(self.__get_general_model_data_only)
         return await document.get_cache(self.__get_general_model)
 
+    async def __get_general_model_data_only(self, document: TextDocument) -> ast.AST:
+        return self.__get_model(document, await self.get_general_tokens(document, True), DocumentType.GENERAL)
+
     async def __get_general_model(self, document: TextDocument) -> ast.AST:
         return self.__get_model(document, await self.get_general_tokens(document), DocumentType.GENERAL)
 
-    async def get_resource_model(self, document: TextDocument) -> ast.AST:
+    async def get_resource_model(self, document: TextDocument, data_only: bool = True) -> ast.AST:
+        if data_only:
+            return await document.get_cache(self.__get_resource_model_data_only)
+
         return await document.get_cache(self.__get_resource_model)
 
+    async def __get_resource_model_data_only(self, document: TextDocument) -> ast.AST:
+        return self.__get_model(document, await self.get_resource_tokens(document, True), DocumentType.RESOURCE)
+
     async def __get_resource_model(self, document: TextDocument) -> ast.AST:
         return self.__get_model(document, await self.get_resource_tokens(document), DocumentType.RESOURCE)
 
-    async def get_init_model(self, document: TextDocument) -> ast.AST:
+    async def get_init_model(self, document: TextDocument, data_only: bool = True) -> ast.AST:
+        if data_only:
+            return await document.get_cache(self.__get_init_model_data_only)
         return await document.get_cache(self.__get_init_model)
 
+    async def __get_init_model_data_only(self, document: TextDocument) -> ast.AST:
+        return self.__get_model(document, await self.get_init_tokens(document, True), DocumentType.INIT)
+
     async def __get_init_model(self, document: TextDocument) -> ast.AST:
         return self.__get_model(document, await self.get_init_tokens(document), DocumentType.INIT)
 
diff --git a/robotcode/language_server/robotframework/parts/semantic_tokens.py b/robotcode/language_server/robotframework/parts/semantic_tokens.py
@@ -726,11 +726,7 @@ async def get_tokens() -> AsyncGenerator[Tuple[Token, ast.AST], None]:
             lambda t: range is None or token_in_range(t[0], range),
             async_dropwhile(
                 lambda t: range is not None and not token_in_range(t[0], range),
-                (
-                    (t, n)
-                    async for t, n in get_tokens()
-                    if t.type not in [RobotToken.SEPARATOR, RobotToken.EOL, RobotToken.EOS]
-                ),
+                ((t, n) async for t, n in get_tokens()),
             ),
         ):
             async for token in self.generate_sem_tokens(
@@ -770,7 +766,7 @@ async def collect_threading(
         self, document: TextDocument, range: Optional[Range]
     ) -> Union[SemanticTokens, SemanticTokensPartialResult, None]:
 
-        model = await self.parent.documents_cache.get_model(document)
+        model = await self.parent.documents_cache.get_model(document, True)
         namespace = await self.parent.documents_cache.get_namespace(document)
 
         builtin_library_doc = next(
diff --git a/robotcode/language_server/robotframework/parts/signature_help.py b/robotcode/language_server/robotframework/parts/signature_help.py
@@ -79,7 +79,7 @@ async def collect(
         self, sender: Any, document: TextDocument, position: Position, context: Optional[SignatureHelpContext] = None
     ) -> Optional[SignatureHelp]:
 
-        result_node = await get_node_at_position(await self.parent.documents_cache.get_model(document), position)
+        result_node = await get_node_at_position(await self.parent.documents_cache.get_model(document, False), position)
         if result_node is None:
             return None
 
diff --git a/robotcode/language_server/robotframework/utils/ast.py b/robotcode/language_server/robotframework/utils/ast.py
@@ -95,6 +95,12 @@ def end_col_offset(self) -> int:
         ...
 
 
+@runtime_checkable
+class HeaderAndBodyBlock(Protocol):
+    header: Any
+    body: List[Any]
+
+
 def range_from_token(token: Token) -> Range:
     return Range(
         start=Position(line=token.lineno - 1, character=token.col_offset),
@@ -105,12 +111,18 @@ def range_from_token(token: Token) -> Range:
     )
 
 
-def range_from_node(node: ast.AST, skip_non_data: bool = False) -> Range:
+def range_from_node(node: ast.AST, skip_non_data: bool = False, only_start: bool = False) -> Range:
     from robot.parsing.lexer import Token as RobotToken
 
     if skip_non_data and isinstance(node, HasTokens) and node.tokens:
         start_token = next((v for v in node.tokens if v.type not in RobotToken.NON_DATA_TOKENS), None)
-        end_token = next((v for v in reversed(node.tokens) if v.type not in RobotToken.NON_DATA_TOKENS), None)
+
+        if only_start and start_token is not None:
+            end_tokens = tuple(t for t in node.tokens if start_token.lineno == t.lineno)
+        else:
+            end_tokens = node.tokens
+
+        end_token = next((v for v in reversed(end_tokens) if v.type not in RobotToken.NON_DATA_TOKENS), None)
         if start_token is not None and end_token is not None:
             return Range(start=range_from_token(start_token).start, end=range_from_token(end_token).end)
 
