Merge branch 'main' of github.com:oraios/serena

pyproject.toml

@@ -59,6 +59,7 @@ dev = [
   "sphinx-toolbox>=3.5.0",
   "sphinxcontrib-bibtex",
   "sphinxcontrib-spelling>=8.0.0",
+  "types-pyyaml>=6.0.12.20241230",
 ]
 
 [project.urls]

src/multilspy/language_server.py

@@ -660,7 +660,9 @@ class LanguageServer:
 
         :param relative_file_path: The relative path of the file that has the symbols
         :param include_body: whether to include the body of the symbols in the result.
-        :return: A list of symbols in the file, and a list of root symbols that represent the tree structure of the symbols. Each symbol in hierarchy starting from the roots has a children attribute.
+        :return: A list of symbols in the file, and a list of root symbols that represent the tree structure of the symbols. 
+            Each symbol in hierarchy starting from the roots has a children attribute.
+            All symbols will have a location and a children attribute.
         """
         self.logger.log(f"Requesting document symbols for {relative_file_path} for the first time", logging.DEBUG)
         # TODO: it's kinda dumb to not use the cache if include_body is False after include_body was True once
@@ -702,7 +704,7 @@ class LanguageServer:
                 item['body'] = self.retrieve_symbol_body(item)
             item[LSPConstants.CHILDREN] = item.get(LSPConstants.CHILDREN, [])
         
-        symbols_without_children: List[multilspy_types.UnifiedSymbolInformation] = []
+        flat_all_symbol_list: List[multilspy_types.UnifiedSymbolInformation] = []
         assert isinstance(response, list)
         root_nodes: List[multilspy_types.UnifiedSymbolInformation] = []
         for item in response:
@@ -722,23 +724,106 @@ class LanguageServer:
                     turn_item_into_symbol_with_children(node)
                     assert LSPConstants.CHILDREN in node
                     children = node[LSPConstants.CHILDREN]
-                    node_without_children = node.copy()
-                    del node_without_children[LSPConstants.CHILDREN]
-                    l.append(node_without_children)
+                    l.append(node)
                     for child in children:
                         l.extend(visit_tree_nodes_and_build_tree_repr(child))
                     return l
                 
-                symbols_without_children.extend(visit_tree_nodes_and_build_tree_repr(item))
+                flat_all_symbol_list.extend(visit_tree_nodes_and_build_tree_repr(item))
             else:
-                symbols_without_children.append(multilspy_types.UnifiedSymbolInformation(**item))
+                flat_all_symbol_list.append(multilspy_types.UnifiedSymbolInformation(**item))
 
-        result = symbols_without_children, root_nodes
+        result = flat_all_symbol_list, root_nodes
         self.logger.log(f"Caching document symbols for {relative_file_path}", logging.DEBUG)
         self._document_symbols_cache[cache_key] = (file_data.content_hash, result)
         self._cache_has_changed = True
         return result
+    
+    async def request_full_symbol_tree(self, start_package_relative_path: str | None = None, include_body: bool = False) -> List[multilspy_types.UnifiedSymbolInformation]:
+        """
+        Will go through all files in the project and build a tree of symbols. Note: this may be slow the first time it is called.
         
+        For each file, a symbol of kind Module (3) will be created. For directories, a symbol of kind Package (4) will be created.
+        All symbols will have a children attribute, thereby representing the tree structure of all symbols in the project 
+        that are within the repository.
+        Will ignore all directories that start with a dot (.) and __pycache__ directories.
+        
+        Args:
+            start_package_relative_path: if passed, only the symbols within this directory will be considered.
+            include_body: whether to include the body of the symbols in the result.
+
+        Returns:
+            A list of root symbols representing the top-level packages/modules in the project.
+        """
+        if not self.server_started:
+            self.logger.log(
+                "request_full_symbol_tree called before Language Server started",
+                logging.ERROR,
+            )
+            raise MultilspyException("Language Server not started")
+
+        # Helper function to check if a path should be ignored
+        def should_ignore_path(path: str) -> bool:
+            parts = path.split(os.sep)
+            return any(part.startswith('.') or part == '__pycache__' for part in parts)
+
+        # Helper function to recursively process directories
+        async def process_directory(dir_path: str) -> List[multilspy_types.UnifiedSymbolInformation]:
+            if should_ignore_path(dir_path):
+                return []
+
+            result = []
+            try:
+                items = os.listdir(os.path.join(self.repository_root_path, dir_path))
+            except OSError:
+                return []
+
+            # Create package symbol for directory
+            package_symbol = multilspy_types.UnifiedSymbolInformation( # type: ignore
+                name=os.path.basename(dir_path),
+                kind=multilspy_types.SymbolKind.Package,
+                location=multilspy_types.Location(
+                    uri=str(pathlib.Path(os.path.join(self.repository_root_path, dir_path)).as_uri()),
+                    range={"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 0}},
+                    absolutePath=str(os.path.join(self.repository_root_path, dir_path)),
+                    relativePath=str(Path(dir_path).resolve().relative_to(self.repository_root_path)),
+                ),
+                children=[]
+            )
+            result.append(package_symbol)
+
+            for item in items:
+                item_path = os.path.join(dir_path, item)
+                abs_item_path = os.path.join(self.repository_root_path, item_path)
+
+                if os.path.isdir(abs_item_path):
+                    child_symbols = await process_directory(item_path)
+                    package_symbol["children"].extend(child_symbols)
+
+                elif os.path.isfile(abs_item_path):
+                    _, root_nodes = await self.request_document_symbols(item_path, include_body=include_body)
+                    
+                    # Create module symbol
+                    module_symbol = multilspy_types.UnifiedSymbolInformation( # type: ignore
+                        name=os.path.splitext(item)[0],
+                        kind=multilspy_types.SymbolKind.Module,
+                        location=multilspy_types.Location(
+                            uri=str(pathlib.Path(abs_item_path).as_uri()),
+                            range={"start": {"line": 0, "character": 0}, "end": {"line": 0, "character": 0}},
+                            absolutePath=str(abs_item_path),
+                            relativePath=str(Path(item_path).resolve().relative_to(self.repository_root_path)),
+                        ),
+                        children=root_nodes
+                    )
+                    
+                    package_symbol["children"].append(module_symbol)
+
+            return result
+
+        # Start from the root or the specified directory
+        start_path = start_package_relative_path or self.repository_root_path
+        return await process_directory(start_path)
+
     async def request_hover(self, relative_file_path: str, line: int, column: int) -> Union[multilspy_types.Hover, None]:
         """
         Raise a [textDocument/hover](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_hover) request to the Language Server
@@ -1190,8 +1275,8 @@ class SyncLanguageServer:
         :return: None
         """
         self.loop = asyncio.new_event_loop()
-        loop_thread = threading.Thread(target=self.loop.run_forever, daemon=True)
-        loop_thread.start()
+        self.loop_thread = threading.Thread(target=self.loop.run_forever, daemon=True)
+        self.loop_thread.start()
         ctx = self.language_server.start_server()
         asyncio.run_coroutine_threadsafe(ctx.__aenter__(), loop=self.loop).result()
         yield self
@@ -1282,6 +1367,28 @@ class SyncLanguageServer:
             self.language_server.request_document_symbols(relative_file_path, include_body), self.loop
         ).result()
         return result
+    
+    def request_full_symbol_tree(self, start_package_relative_path: str | None = None, include_body: bool = False) -> List[multilspy_types.UnifiedSymbolInformation]:
+        """
+        Will go through all files in the project and build a tree of symbols. Note: this may be slow the first time it is called.
+        
+        For each file, a symbol of kind Module (3) will be created. For directories, a symbol of kind Package (4) will be created.
+        All symbols will have a children attribute, thereby representing the tree structure of all symbols in the project 
+        that are within the repository.
+        Will ignore all directories that start with a dot (.) and __pycache__ directories.
+        
+        Args:
+            start_package_relative_path: if passed, only the symbols within this directory will be considered.
+            include_body: whether to include the body of the symbols in the result.
+
+        Returns:
+            A list of root symbols representing the top-level packages/modules in the project.
+        """
+        result = asyncio.run_coroutine_threadsafe(
+            self.language_server.request_full_symbol_tree(start_package_relative_path, include_body), self.loop
+        ).result()
+        return result
+
 
     def request_hover(self, relative_file_path: str, line: int, column: int) -> Union[multilspy_types.Hover, None]:
         """

src/multilspy/multilspy_types.py

@@ -215,14 +215,12 @@ class UnifiedSymbolInformation(TypedDict):
     body: NotRequired[str]
     """ The body of the symbol. """
     
-    children: NotRequired[List[UnifiedSymbolInformation]]
+    children: List[UnifiedSymbolInformation]
     """ The children of the symbol. 
     Added to be compatible with `lsp_types.DocumentSymbol`, 
     since it is sometimes useful to have the children of the symbol as a user-facing feature."""
     
 
-TreeRepr = Dict[int, List['TreeRepr']]
-
 class MarkupKind(Enum):
     """Describes the content type that a client supports in various
     result literals like `Hover`, `ParameterInfo` or `CompletionItem`.

src/serena/llm/jinja_template.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import jinja2
 import jinja2.meta
 import jinja2.nodes
@@ -8,21 +10,21 @@ from serena.util.class_decorators import singleton
 
 @singleton
 class JinjaEnvProvider:
-    def __init__(self):
-        self._env = None
+    def __init__(self) -> None:
+        self._env: jinja2.Environment | None = None
 
-    def get_env(self):
+    def get_env(self) -> jinja2.Environment:
         if self._env is None:
             self._env = jinja2.Environment()
         return self._env
 
 
 class JinjaTemplate:
-    def __init__(self, template_string: str):
+    def __init__(self, template_string: str) -> None:
         self._template_string = template_string
         self._template = JinjaEnvProvider().get_env().from_string(self._template_string)
 
-    def render(self, **kwargs) -> str:
+    def render(self, **kwargs: Any) -> str:
         return self._template.render(**kwargs)
 
     def get_parameters(self) -> set[str]:

src/serena/llm/multilang_prompt.py

@@ -12,7 +12,7 @@ LANG_CODES = ["en", "de"]
 
 
 class PromptTemplate(ToStringMixin):
-    def __init__(self, name: str, jinja_template_string: str):
+    def __init__(self, name: str, jinja_template_string: str) -> None:
         self.name = name
         self.jinja_template = JinjaTemplate(jinja_template_string.strip())
         self.parameters = self.jinja_template.get_parameters()
@@ -20,15 +20,15 @@ class PromptTemplate(ToStringMixin):
     def _tostring_excludes(self) -> list[str]:
         return ["jinja_template"]
 
-    def instantiate(self, **kwargs) -> str:
+    def instantiate(self, **kwargs: Any) -> str:
         return self.jinja_template.render(**kwargs)
 
 
 class PromptList:
-    def __init__(self, items: list[str]):
+    def __init__(self, items: list[str]) -> None:
         self.items = [x.strip() for x in items]
 
-    def to_string(self):
+    def to_string(self) -> str:
         bullet = " * "
         indent = " " * len(bullet)
         items = [x.replace("\n", "\n" + indent) for x in self.items]
@@ -43,7 +43,7 @@ class MultiLangContainer(Generic[T], ToStringMixin):
     Represents a container of items which are associated with different languages
     """
 
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         self.name = name
         self.lang2item: dict[str, T] = {}
 
@@ -63,7 +63,7 @@ class MultiLangContainer(Generic[T], ToStringMixin):
         If the requested language is not found, raise an exception
         """
 
-    def add_item(self, item: T, lang: str = ""):
+    def add_item(self, item: T, lang: str = "") -> None:
         self.lang2item[lang] = item
 
     def get_item(self, lang: str, fallback_mode: FallbackMode = FallbackMode.EXCEPTION) -> T:
@@ -105,6 +105,7 @@ class MultiLangPromptTemplate(MultiLangContainer[PromptTemplate]):
                     params == prev_params
                 ), f"Parameters of MLPT '{self.name}' are inconsistent: {sorted(params)} vs {sorted(prev_params)}"
             prev_params = params
+        assert prev_params is not None
         return sorted(prev_params)
 
 
@@ -126,7 +127,7 @@ class MultiLangPromptTemplateCollection:
           The language of all can be set by specifying the key 'lang' in addition to 'prompts'.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.prompt_templates: dict[str, MultiLangPromptTemplate] = {}
         self.prompt_lists: dict[str, MultiLangPromptList] = {}
         prompts_dir = self._prompt_template_folder()
@@ -141,7 +142,7 @@ class MultiLangPromptTemplateCollection:
             prompts_dir = os.path.join(dir_path, "prompts")
             if os.path.isdir(prompts_dir):
                 break
-        if not os.path.isdir(prompts_dir):
+        if prompts_dir is None or not os.path.isdir(prompts_dir):
             raise FileNotFoundError("Could not find the 'prompts' directory")
         return prompts_dir
 
@@ -159,7 +160,7 @@ class MultiLangPromptTemplateCollection:
 
         return container, lang
 
-    def _add_prompt_template(self, prompt_name: str, jinja_prompt_template: str):
+    def _add_prompt_template(self, prompt_name: str, jinja_prompt_template: str) -> None:
         """
         :param prompt_name: a prompt name, which may have a language shortcode suffix (e.g. "_de")
         :param jinja_prompt_template: the actual prompt string which may contain placeholders/parameters (e.g. "{name}")
@@ -167,7 +168,7 @@ class MultiLangPromptTemplateCollection:
         multilang_prompt_template, lang = self._container_lang(prompt_name, self.prompt_templates, MultiLangPromptTemplate)
         multilang_prompt_template.add_item(PromptTemplate(prompt_name, jinja_prompt_template), lang=lang)
 
-    def _add_prompt_list(self, prompt_name: str, prompt_list: list[str]):
+    def _add_prompt_list(self, prompt_name: str, prompt_list: list[str]) -> None:
         """
         :param prompt_name: a prompt name, which may have a language shortcode suffix (e.g. "_de")
         :param prompt_list: a list of prompts
@@ -175,7 +176,7 @@ class MultiLangPromptTemplateCollection:
         multilang_prompt_list, lang = self._container_lang(prompt_name, self.prompt_lists, MultiLangPromptList)
         multilang_prompt_list.add_item(PromptList(prompt_list), lang=lang)
 
-    def _read_prompt_templates(self, prompts_dir: str):
+    def _read_prompt_templates(self, prompts_dir: str) -> None:
         for fn in os.listdir(prompts_dir):
             path = os.path.join(prompts_dir, fn)
             if fn.endswith(".txt"):

src/serena/llm/prompt_factory.py

@@ -4,12 +4,14 @@ from .multilang_prompt import MultiLangContainer, MultiLangPromptTemplateCollect
 class PromptFactory:
     # NOTE: This class is auto-generated by gen_prompt_factory.py
 
-    def __init__(self, lang_shortcode: str = "en", fallback_mode=MultiLangContainer.FallbackMode.EXCEPTION):
+    def __init__(
+        self, lang_shortcode: str = "en", fallback_mode: MultiLangContainer.FallbackMode = MultiLangContainer.FallbackMode.EXCEPTION
+    ):
         self.lang_shortcode = lang_shortcode
         self.collection = MultiLangPromptTemplateCollection()
         self.fallback_mode = fallback_mode
 
-    def _format_prompt(self, prompt_name: str, kwargs) -> str:
+    def _format_prompt(self, prompt_name: str, kwargs: dict) -> str:
         del kwargs["self"]
         mpt = self.collection.get_multilang_prompt_template(prompt_name)
         return mpt.get_item(self.lang_shortcode, self.fallback_mode).instantiate(**kwargs)
@@ -18,5 +20,5 @@ class PromptFactory:
         mpl = self.collection.get_multilang_prompt_list(prompt_name)
         return mpl.get_item(self.lang_shortcode, self.fallback_mode)
 
-    def create_onboarding_prompt(self, *, onboarding_file) -> str:
+    def create_onboarding_prompt(self, *, onboarding_file: str) -> str:
         return self._format_prompt("onboarding_prompt", locals())

src/serena/text_utils.py

@@ -52,7 +52,7 @@ class MatchedConsecutiveLines:
     matched_lines: list[TextLine] = field(default_factory=list)
     lines_after_matched: list[TextLine] = field(default_factory=list)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         for line in self.lines:
             if line.match_type == LineType.BEFORE_MATCH:
                 self.lines_before_matched.append(line)

src/serena/util/class_decorators.py

@@ -1,7 +1,10 @@
-def singleton(cls):
+from typing import Any
+
+
+def singleton(cls: type[Any]) -> Any:
     instance = None
 
-    def get_instance(*args, **kwargs):
+    def get_instance(*args: Any, **kwargs: Any) -> Any:
         nonlocal instance
         if instance is None:
             instance = cls(*args, **kwargs)

src/serena/util/file_system.py

@@ -4,7 +4,7 @@ from collections.abc import Sequence
 
 def scan_directory(
     path: str,
-    recursive=False,
+    recursive: bool = False,
     relative_to: str | None = None,
     ignored_dirs: Sequence[str] = (),
     ignored_files: Sequence[str] = (),
@@ -24,7 +24,7 @@ def scan_directory(
     rel_base = os.path.abspath(relative_to) if relative_to else None
 
     # Helper function to check if an item should be ignored
-    def is_ignored(entry_path, ignored_items):
+    def is_ignored(entry_path: str, ignored_items: Sequence[str]) -> bool:
         entry_name = os.path.basename(entry_path)
 
         # Check if name is directly in ignored list

test/multilspy/test_symbol_retrieval.py

@@ -293,6 +293,8 @@ class TestLanguageServerSymbols:
 
         # Step 3: Verify that they refer to the same symbol
         assert defining_symbol["kind"] == containing_symbol["kind"]
+        assert "location" in defining_symbol
+        assert "location" in containing_symbol
         assert defining_symbol["location"]["uri"] == containing_symbol["location"]["uri"]
 
         # The integration test is successful if we've gotten this far,
@@ -320,35 +322,46 @@ class TestLanguageServerSymbols:
                 warnings.warn("Could not verify container hierarchy - implementation detail")
 
     def test_symbol_tree_structure(self, language_server: SyncLanguageServer, repo_path: Path):
-        """Test the symbol tree structure."""
-        file_path = str(repo_path / "test_repo" / "services.py")
-        symbols, root_nodes = language_server.request_document_symbols(file_path)
-        assert len(symbols) > 0
-        assert {root["name"] for root in root_nodes} == {
-            "UserService",
-            "ItemService",
-            "create_service_container",
-            "user_var_str",
-            "user_service",
-        }
-        user_service_root = next(root for root in root_nodes if root["name"] == "UserService")
-        assert user_service_root
-        assert "children" in user_service_root
-        assert {child["name"] for child in user_service_root["children"] if child["kind"] != SymbolKind.Variable} == {
-            "__init__",
-            "create_user",
-            "get_user",
-            "list_users",
-            "delete_user",
-        }
+        """Test that the symbol tree structure is correctly built."""
+        # Get all symbols in the test file
+        repo_structure = language_server.request_full_symbol_tree()
+        assert len(repo_structure) == 1
+        # Assert that the root symbol is the test_repo directory
+        assert repo_structure[0]["name"] == "test_repo"
+        assert repo_structure[0]["kind"] == SymbolKind.Package
+        assert "children" in repo_structure[0]
+        # Assert that the children are the top-level packages
+        child_names = {child["name"] for child in repo_structure[0]["children"]}
+        child_kinds = {child["kind"] for child in repo_structure[0]["children"]}
+        assert child_names == {"test_repo", "custom_test", "examples", "scripts"}
+        assert child_kinds == {SymbolKind.Package}
+        examples_package = next(child for child in repo_structure[0]["children"] if child["name"] == "examples")
+        # assert that children are __init__ and user_management
+        assert {child["name"] for child in examples_package["children"]} == {"__init__", "user_management"}
+        assert {child["kind"] for child in examples_package["children"]} == {SymbolKind.Module}
 
-        # Now recursively flatten the tree into a set of names and assert that it coincides with the set of symbol names
-        all_names_in_tree = set()
+        # assert that tree of user_management node is same as retrieved directly
+        user_management_node = next(child for child in examples_package["children"] if child["name"] == "user_management")
+        user_management_rel_path = user_management_node["location"]["relativePath"]
+        assert user_management_rel_path == "examples/user_management.py"
+        _, user_management_roots = language_server.request_document_symbols(str(repo_path / "examples" / "user_management.py"))
+        assert user_management_roots == user_management_node["children"]
 
-        def flatten_tree(nodes):
-            for node in nodes:
-                all_names_in_tree.add(node["name"])
-                flatten_tree(node.get("children", []))
+    def test_symbol_tree_structure_subdir(self, language_server: SyncLanguageServer, repo_path: Path):
+        """Test that the symbol tree structure is correctly built."""
+        # Get all symbols in the test file
+        examples_package_roots = language_server.request_full_symbol_tree(start_package_relative_path=str(repo_path / "examples"))
+        assert len(examples_package_roots) == 1
+        examples_package = examples_package_roots[0]
+        assert examples_package["name"] == "examples"
+        assert examples_package["kind"] == SymbolKind.Package
+        # assert that children are __init__ and user_management
+        assert {child["name"] for child in examples_package["children"]} == {"__init__", "user_management"}
+        assert {child["kind"] for child in examples_package["children"]} == {SymbolKind.Module}
 
-        flatten_tree(root_nodes)
-        assert all_names_in_tree == {symbol["name"] for symbol in symbols}
+        # assert that tree of user_management node is same as retrieved directly
+        user_management_node = next(child for child in examples_package["children"] if child["name"] == "user_management")
+        user_management_rel_path = user_management_node["location"]["relativePath"]
+        assert user_management_rel_path == "examples/user_management.py"
+        _, user_management_roots = language_server.request_document_symbols(str(repo_path / "examples" / "user_management.py"))
+        assert user_management_roots == user_management_node["children"]

uv.lock

@@ -2114,6 +2114,7 @@ dev = [
     { name = "sphinxcontrib-bibtex" },
     { name = "sphinxcontrib-spelling" },
     { name = "toml-sort" },
+    { name = "types-pyyaml" },
 ]
 
 [package.metadata]
@@ -2151,6 +2152,7 @@ requires-dist = [
     { name = "sphinxcontrib-bibtex", marker = "extra == 'dev'" },
     { name = "sphinxcontrib-spelling", marker = "extra == 'dev'", specifier = ">=8.0.0" },
     { name = "toml-sort", marker = "extra == 'dev'", specifier = ">=0.24.2" },
+    { name = "types-pyyaml", marker = "extra == 'dev'", specifier = ">=6.0.12.20241230" },
 ]
 provides-extras = ["dev"]
 
@@ -2709,6 +2711,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0f/b3/ca41df24db5eb99b00d97f89d7674a90cb6b3134c52fb8121b6d8d30f15c/types_python_dateutil-2.9.0.20241206-py3-none-any.whl", hash = "sha256:e248a4bc70a486d3e3ec84d0dc30eec3a5f979d6e7ee4123ae043eedbb987f53", size = 14384 },
 ]
 
+[[package]]
+name = "types-pyyaml"
+version = "6.0.12.20241230"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/9a/f9/4d566925bcf9396136c0a2e5dc7e230ff08d86fa011a69888dd184469d80/types_pyyaml-6.0.12.20241230.tar.gz", hash = "sha256:7f07622dbd34bb9c8b264fe860a17e0efcad00d50b5f27e93984909d9363498c", size = 17078 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e8/c1/48474fbead512b70ccdb4f81ba5eb4a58f69d100ba19f17c92c0c4f50ae6/types_PyYAML-6.0.12.20241230-py3-none-any.whl", hash = "sha256:fa4d32565219b68e6dee5f67534c722e53c00d1cfc09c435ef04d7353e1e96e6", size = 20029 },
+]
+
 [[package]]
 name = "typing-extensions"
 version = "4.12.2"