speedup semantic token, if installed use orjson package to speedup json de/encoding

Author: d-biehl

package.json

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

pyproject.toml

@@ -28,10 +28,11 @@ robotremoteserver = "^1.1"
 Cython = "^0.29.24"
 robotframework-robocop = "^1.7.1"
 robotframework-tidy = "^1.5.1"
+orjson = "^3.6.3"
 
 
 [tool.poetry-dynamic-versioning]
-enable = true
+enable = false
 vcs = "git"
 style = "semver"
 

robotcode/jsonrpc2/protocol.py

@@ -75,9 +75,46 @@ class JsonRPCErrors:
 PROTOCOL_VERSION = "2.0"
 
 
+def get_loads() -> Callable[..., Any]:
+
+    try:
+        import orjson
+
+        return cast(Callable[..., Any], orjson.loads)
+
+    except ImportError:
+        pass
+
+    import json
+
+    return json.loads
+
+
+def get_dumps() -> Callable[..., Any]:
+    try:
+        import orjson
+
+        def orjson_dumps(__obj: Any, *, default: Any, indent: Optional[bool] = False) -> str:
+            # orjson.dumps returns bytes, to match standard json.dumps we need to decode
+            return orjson.dumps(__obj, default=default, option=orjson.OPT_INDENT_2 if indent else 0).decode()
+
+        return orjson_dumps
+
+    except ImportError:
+        pass
+
+    import json
+
+    return json.dumps
+
+
 class JsonRPCMessage(BaseModel):
     jsonrpc: str = Field(PROTOCOL_VERSION, const=True)
 
+    class Config:
+        json_loads = get_loads()
+        json_dumps = get_dumps()
+
 
 class JsonRPCNotification(JsonRPCMessage):
     method: str = Field(...)

robotcode/language_server/robotframework/parts/semantic_tokens.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import itertools
 import operator
 import re
 from enum import Enum
@@ -96,6 +97,7 @@ class RobotSemanticTokenProtocolPart(RobotLanguageServerProtocolPart):
     def __init__(self, parent: RobotLanguageServerProtocol) -> None:
         super().__init__(parent)
         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)
@@ -188,7 +190,7 @@ def generate_sem_sub_tokens(
                 else:
                     yield SemTokenInfo.from_token(token, sem_info[0], sem_info[1])
 
-            if token.type == RobotToken.ARGUMENT and "\\" in token.value:
+            elif token.type == RobotToken.ARGUMENT and "\\" in token.value:
                 if col_offset is None:
                     col_offset = token.col_offset
                 if length is None:
@@ -241,19 +243,10 @@ def collect(
         last_line = 0
         last_col = 0
 
-        start = True
-
-        for robot_token in tokens:
-            if range is not None:
-                if start:
-                    if not token_in_range(robot_token, range):
-                        continue
-                    else:
-                        start = False
-                else:
-                    if not token_in_range(robot_token, range):
-                        break
-
+        for robot_token in itertools.takewhile(
+            lambda t: range is not None and token_in_range(t, range),
+            itertools.dropwhile(lambda t: range is None or not token_in_range(t, range), tokens),
+        ):
             for token in self.generate_sem_tokens(robot_token):
                 current_line = token.lineno - 1