find references, highlight, rename tags

d-biehl · d-biehl · commit b1401224e304 · 2022-05-07T12:04:59.000+02:00
diff --git a/robotcode/jsonrpc2/protocol.py b/robotcode/jsonrpc2/protocol.py
@@ -581,7 +581,11 @@ def handle_response(self, message: JsonRPCResponse) -> None:
                 if entry.future._loop == asyncio.get_running_loop():
                     entry.future.set_result(res)
                 else:
-                    entry.future._loop.call_soon_threadsafe(entry.future.set_result, res)
+                    if entry.future._loop.is_running():
+                        entry.future._loop.call_soon_threadsafe(entry.future.set_result, res)
+                    else:
+                        self._logger.warning("Response loop is not running.")
+
         except (SystemExit, KeyboardInterrupt):
             raise
         except BaseException as e:
diff --git a/robotcode/language_server/common/parts/code_lens.py b/robotcode/language_server/common/parts/code_lens.py
@@ -93,4 +93,4 @@ async def refresh(self) -> None:
         ):
             return
 
-        await self.parent.send_request_async("workspace/codeLens/refresh")
+        return await self.parent.send_request("workspace/codeLens/refresh")
diff --git a/robotcode/language_server/robotframework/parts/codelens.py b/robotcode/language_server/robotframework/parts/codelens.py
@@ -3,7 +3,7 @@
 import ast
 from typing import TYPE_CHECKING, Any, List, Optional, cast
 
-from ....utils.async_tools import threaded
+from ....utils.async_tools import create_sub_task, threaded
 from ....utils.logging import LoggingDescriptor
 from ...common.decorators import language_id
 from ...common.lsp_types import CodeLens, Command
@@ -105,14 +105,35 @@ async def resolve(self, sender: Any, code_lens: CodeLens) -> Optional[CodeLens]:
             kw_doc = await self.get_keyword_definition_at_line(namespace, name, line)
 
             if kw_doc is not None and not kw_doc.is_error_handler:
-                references = await self.parent.robot_references.find_keyword_references(
+                if not await self.parent.robot_references.has_cached_keyword_references(
                     document, kw_doc, include_declaration=False
-                )
-                code_lens.command = Command(
-                    f"{len(references)} references",
-                    "editor.action.showReferences",
-                    [str(document.uri), code_lens.range.start, references],
-                )
+                ):
+                    code_lens.command = Command(
+                        "...",
+                        "editor.action.showReferences",
+                        [str(document.uri), code_lens.range.start, []],
+                    )
+
+                    async def find_refs() -> None:
+                        if document is None or kw_doc is None:
+                            return
+
+                        await self.parent.robot_references.find_keyword_references(
+                            document, kw_doc, include_declaration=False
+                        )
+
+                        await self.parent.code_lens.refresh()
+
+                    create_sub_task(find_refs(), loop=self.parent.loop)
+                else:
+                    references = await self.parent.robot_references.find_keyword_references(
+                        document, kw_doc, include_declaration=False
+                    )
+                    code_lens.command = Command(
+                        f"{len(references)} references",
+                        "editor.action.showReferences",
+                        [str(document.uri), code_lens.range.start, references],
+                    )
             else:
                 code_lens.command = Command(
                     "0 references",
diff --git a/robotcode/language_server/robotframework/parts/document_highlight.py b/robotcode/language_server/robotframework/parts/document_highlight.py
@@ -332,3 +332,39 @@ async def highlight_Template(  # noqa: N802
         self, result_node: ast.AST, document: TextDocument, position: Position
     ) -> Optional[List[DocumentHighlight]]:
         return await self._highlight_Template_or_TestTemplate(result_node, document, position)
+
+    async def _highlight_ForceTags_DefaultTags_Tags(  # noqa: N802
+        self, result_node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[List[DocumentHighlight]]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+        from robot.utils.normalizing import normalize
+
+        token = get_tokens_at_position(cast(HasTokens, result_node), position)[-1]
+
+        if token is None:
+            return None
+
+        if token.type in [RobotToken.ARGUMENT] and token.value:
+            return [
+                DocumentHighlight(e.range, DocumentHighlightKind.TEXT)
+                for e in await self.parent.robot_references.find_tag_references_in_file(
+                    document, normalize(token.value, ignore="_")
+                )
+            ]
+
+        return None
+
+    async def highlight_ForceTags(  # noqa: N802
+        self, result_node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[List[DocumentHighlight]]:
+        return await self._highlight_ForceTags_DefaultTags_Tags(result_node, document, position)
+
+    async def highlight_DefaultTags(  # noqa: N802
+        self, result_node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[List[DocumentHighlight]]:
+        return await self._highlight_ForceTags_DefaultTags_Tags(result_node, document, position)
+
+    async def highlight_Tags(  # noqa: N802
+        self, result_node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[List[DocumentHighlight]]:
+        return await self._highlight_ForceTags_DefaultTags_Tags(result_node, document, position)
diff --git a/robotcode/language_server/robotframework/parts/references.py b/robotcode/language_server/robotframework/parts/references.py
@@ -18,6 +18,7 @@
     cast,
 )
 
+from ....utils.async_cache import AsyncSimpleCache
 from ....utils.async_itertools import async_next
 from ....utils.async_tools import create_sub_task, run_coroutine_in_thread, threaded
 from ....utils.glob_path import iter_files
@@ -71,6 +72,8 @@ def __init__(self, parent: RobotLanguageServerProtocol) -> None:
 
         parent.references.collect.add(self.collect)
 
+        self._keyword_reference_cache = AsyncSimpleCache(max_items=128)
+
     def _find_method(self, cls: Type[Any]) -> Optional[_ReferencesMethod]:
         if cls is ast.AST:
             return None
@@ -733,9 +736,21 @@ async def get_keyword_references_from_any_run_keyword(
                     async for e in self.get_keyword_references_from_tokens(namespace, kw_doc, node, t, args, True):
                         yield e
 
+    async def has_cached_keyword_references(
+        self, document: TextDocument, kw_doc: KeywordDoc, include_declaration: bool = True
+    ) -> bool:
+        return await self._keyword_reference_cache.has(document, kw_doc, include_declaration)
+
     async def find_keyword_references(
         self, document: TextDocument, kw_doc: KeywordDoc, include_declaration: bool = True
     ) -> List[Location]:
+        return await self._keyword_reference_cache.get(
+            self._find_keyword_references, document, kw_doc, include_declaration
+        )
+
+    async def _find_keyword_references(
+        self, document: TextDocument, kw_doc: KeywordDoc, include_declaration: bool = True
+    ) -> List[Location]:
 
         namespace = await self.parent.documents_cache.get_namespace(document)
         if namespace is None:
@@ -923,3 +938,63 @@ async def references_VariablesImport(  # noqa: N802
             )
 
         return None
+
+    async def find_tag_references_in_file(
+        self, doc: TextDocument, tag: str, is_normalized: bool = False
+    ) -> List[Location]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+        from robot.parsing.model.statements import DefaultTags, ForceTags, Tags
+        from robot.utils.normalizing import normalize
+
+        model = await self.parent.documents_cache.get_model(doc)
+        if model is None:
+            return []
+
+        result: List[Location] = []
+        if not is_normalized:
+            tag = normalize(tag, ignore="_")
+
+        async for node in iter_nodes(model):
+            if isinstance(node, (ForceTags, DefaultTags, Tags)):
+                for token in node.get_tokens(RobotToken.ARGUMENT):
+                    if token.value and normalize(token.value, ignore="_") == tag:
+                        result.append(Location(str(doc.uri), range_from_token(token)))
+
+        return result
+
+    async def _references_ForceTags_DefaultTags_Tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, context: ReferenceContext
+    ) -> Optional[List[Location]]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+
+        token = get_tokens_at_position(cast(HasTokens, node), position)[-1]
+
+        if token is None:
+            return None
+
+        if token.type in [RobotToken.ARGUMENT] and token.value:
+            return await self.find_tag_references(document, token.value)
+
+        return None
+
+    async def find_tag_references(self, document: TextDocument, tag: str) -> List[Location]:
+        from robot.utils.normalizing import normalize
+
+        return await self._find_references_in_workspace(
+            document, self.find_tag_references_in_file, normalize(tag, ignore="_"), True
+        )
+
+    async def references_ForceTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, context: ReferenceContext
+    ) -> Optional[List[Location]]:
+        return await self._references_ForceTags_DefaultTags_Tags(node, document, position, context)
+
+    async def references_DefaultTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, context: ReferenceContext
+    ) -> Optional[List[Location]]:
+        return await self._references_ForceTags_DefaultTags_Tags(node, document, position, context)
+
+    async def references_Tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, context: ReferenceContext
+    ) -> Optional[List[Location]]:
+        return await self._references_ForceTags_DefaultTags_Tags(node, document, position, context)
diff --git a/robotcode/language_server/robotframework/parts/rename.py b/robotcode/language_server/robotframework/parts/rename.py
@@ -550,3 +550,79 @@ async def rename_Template(  # noqa: N802
         return await self._rename_keyword(
             document, new_name, await self._find_Template_or_TestTemplate(node, document, position)
         )
+
+    async def _prepare_rename_tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[PrepareRenameResult]:
+        from robot.parsing.lexer.tokens import Token as RobotToken
+
+        token = get_tokens_at_position(cast(HasTokens, node), position)[-1]
+
+        if token is None:
+            return None
+
+        if token.type in [RobotToken.ARGUMENT] and token.value:
+            return PrepareRenameResultWithPlaceHolder(range_from_token(token), token.value)
+
+        return None
+
+    async def prepare_rename_ForceTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[PrepareRenameResult]:
+        return await self._prepare_rename_tags(node, document, position)
+
+    async def prepare_rename_DefaultTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[PrepareRenameResult]:
+        return await self._prepare_rename_tags(node, document, position)
+
+    async def prepare_rename_Tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position
+    ) -> Optional[PrepareRenameResult]:
+        return await self._prepare_rename_tags(node, document, position)
+
+    async def _rename_tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, new_name: str
+    ) -> Optional[WorkspaceEdit]:
+
+        from robot.parsing.lexer.tokens import Token as RobotToken
+
+        token = get_tokens_at_position(cast(HasTokens, node), position)[-1]
+
+        if token is None:
+            return None
+
+        if token.type in [RobotToken.ARGUMENT] and token.value:
+            references = await self.parent.robot_references.find_tag_references(document, token.value)
+
+            changes: List[Union[TextDocumentEdit, CreateFile, RenameFile, DeleteFile]] = []
+
+            for reference in references:
+                changes.append(
+                    TextDocumentEdit(
+                        OptionalVersionedTextDocumentIdentifier(reference.uri, None),
+                        [AnnotatedTextEdit(reference.range, new_name, annotation_id="rename_tag")],
+                    )
+                )
+
+            return WorkspaceEdit(
+                document_changes=changes,
+                change_annotations={"rename_keyword": ChangeAnnotation("Rename Tag", False)},
+            )
+
+        return None
+
+    async def rename_ForceTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, new_name: str
+    ) -> Optional[WorkspaceEdit]:
+        return await self._rename_tags(node, document, position, new_name)
+
+    async def rename_DefaultTags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, new_name: str
+    ) -> Optional[WorkspaceEdit]:
+        return await self._rename_tags(node, document, position, new_name)
+
+    async def rename_Tags(  # noqa: N802
+        self, node: ast.AST, document: TextDocument, position: Position, new_name: str
+    ) -> Optional[WorkspaceEdit]:
+        return await self._rename_tags(node, document, position, new_name)
diff --git a/robotcode/language_server/robotframework/parts/robot_workspace.py b/robotcode/language_server/robotframework/parts/robot_workspace.py
@@ -102,6 +102,7 @@ async def _collect_workspace_documents(self, sender: Any) -> List[WorkspaceDocum
                         self._logger.exception(e)
 
         self.workspace_loaded = True
+
         if config.analysis.references_code_lens:
             await self.parent.code_lens.refresh()
 
diff --git a/robotcode/utils/async_cache.py b/robotcode/utils/async_cache.py
@@ -1,3 +1,5 @@
+from collections import defaultdict
+from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Dict, List, Tuple, TypeVar, cast
 
 from .async_tools import Lock
@@ -7,37 +9,47 @@
 
 def _freeze(v: Any) -> Any:
     if isinstance(v, dict):
-        return frozenset(v.items())
+        return hash(frozenset(v.items()))
     return v
 
 
+@dataclass
+class CacheEntry:
+    value: Any = None
+    has_value: bool = False
+    lock: Lock = Lock()
+
+
 class AsyncSimpleCache:
     def __init__(self, max_items: int = 128) -> None:
         self.max_items = max_items
 
-        self._cache: Dict[Tuple[Any, ...], Any] = {}
+        self._cache: Dict[Tuple[Any, ...], CacheEntry] = defaultdict(CacheEntry)
         self._order: List[Tuple[Any, ...]] = []
         self._lock = Lock()
 
+    async def has(self, *args: Any, **kwargs: Any) -> bool:
+        return self._make_key(*args, **kwargs) in self._cache
+
     async def get(self, func: Callable[..., Awaitable[_T]], *args: Any, **kwargs: Any) -> _T:
         key = self._make_key(*args, **kwargs)
 
         async with self._lock:
-            try:
-                return cast(_T, self._cache[key])
-            except KeyError:
-                pass
+            entry = self._cache[key]
+            if entry.has_value:
+                return cast(_T, entry.value)
 
-            res = await func(*args, **kwargs)
+        async with entry.lock:
+            entry.value = await func(*args, **kwargs)
+            entry.has_value = True
 
-            self._cache[key] = res
             self._order.insert(0, key)
 
             if len(self._order) > self.max_items:
                 del self._cache[self._order.pop()]
 
-            return res
+            return entry.value
 
     @staticmethod
     def _make_key(*args: Any, **kwargs: Any) -> Tuple[Any, ...]:
-        return (tuple(_freeze(v) for v in args), frozenset({k: _freeze(v) for k, v in kwargs.items()}))
+        return (tuple(_freeze(v) for v in args), hash(frozenset({k: _freeze(v) for k, v in kwargs.items()})))
diff --git a/tests/robotcode/language_server/robotframework/parts/data/tests/__init__.robot b/tests/robotcode/language_server/robotframework/parts/data/tests/__init__.robot
@@ -1,4 +1,4 @@
 *** Settings ***
 Resource    firstresource.resource
 
-Suite Setup   do something in a resource
+Suite Setup   do something in a resource
diff --git a/tests/robotcode/language_server/robotframework/parts/data/tests/tags/first.robot b/tests/robotcode/language_server/robotframework/parts/data/tests/tags/first.robot
@@ -0,0 +1,12 @@
+*** Settings ***
+Default Tags    blah-tag    bluf-tag
+
+
+*** Test Cases ***
+first
+    [Tags]    no-ci_1
+    No Operation
+
+second
+    [Tags]    unknown 1
+    No Operation
diff --git a/tests/robotcode/language_server/robotframework/parts/data/tests/tags/second.robot b/tests/robotcode/language_server/robotframework/parts/data/tests/tags/second.robot
@@ -0,0 +1,12 @@
+*** Settings ***
+Default Tags    blah-tag    bluf-tag
+
+
+*** Test Cases ***
+first
+    [Tags]    unknown
+    No Operation
+
+second
+    [Tags]    unknown    no-ci_1
+    No Operation
diff --git a/tests/robotcode/language_server/robotframework/parts/data/tests/tags/sub/__init__.robot b/tests/robotcode/language_server/robotframework/parts/data/tests/tags/sub/__init__.robot
diff --git a/tests/robotcode/language_server/robotframework/parts/data/tests/tags/sub/first.robot b/tests/robotcode/language_server/robotframework/parts/data/tests/tags/sub/first.robot
