refactor: better support for multiple workspace folders

Previously, the auto venv detection only works for
the first workspace folder and all workspace folder
will use that detection result.

Now the auto venv detection works per workspace folder.

Signed-off-by: Jack Cherng <jfcherng@gmail.com>

plugin/client.py

@@ -3,58 +3,43 @@
 import json
 import os
 import re
-import shutil
-import weakref
-from dataclasses import dataclass
-from pathlib import Path
 from typing import Any, cast
 
 import jmespath
 import sublime
+import sublime_plugin
 from LSP.plugin import ClientConfig, DottedDict, MarkdownLangMap, Response, WorkspaceFolder
 from LSP.plugin.core.protocol import CompletionItem, Hover, SignatureHelp
 from lsp_utils import NpmClientHandler
-from more_itertools import first_true
 from sublime_lib import ResourcePath
 
 from .constants import PACKAGE_NAME, SERVER_SETTING_DEV_ENVIRONMENT
 from .dev_environment.helpers import get_dev_environment_handler
 from .log import log_error, log_info, log_warning
-from .template import load_string_template
-from .virtual_env.helpers import find_venv_by_finder_names, find_venv_by_python_executable
-from .virtual_env.venv_finder import BaseVenvInfo, get_finder_name_mapping
+from .utils_lsp import AbstractLspPythonPlugin, find_workspace_folder, update_view_status_bar_text, uri_to_file_path
+from .virtual_env.helpers import find_venv_by_finder_names
 
 
-@dataclass
-class WindowAttr:
-    simple_python_executable: Path | None = None
-    """The path to the Python executable found by the `PATH` env variable."""
-    venv_info: BaseVenvInfo | None = None
-    """The information of the virtual environment."""
+class ViewEventListener(sublime_plugin.ViewEventListener):
+    def on_activated(self) -> None:
+        settings = self.view.settings()
 
-    @property
-    def preferred_python_executable(self) -> Path | None:
-        return self.venv_info.python_executable if self.venv_info else self.simple_python_executable
+        if settings.get("lsp_active"):
+            update_view_status_bar_text(LspPyrightPlugin, self.view)
 
 
-class LspPyrightPlugin(NpmClientHandler):
+class LspPyrightPlugin(AbstractLspPythonPlugin, NpmClientHandler):
     package_name = PACKAGE_NAME
     server_directory = "language-server"
     server_binary_path = os.path.join(server_directory, "node_modules", "pyright", "langserver.index.js")
 
-    server_version = ""
-    """The version of the language server."""
-
-    window_attrs: weakref.WeakKeyDictionary[sublime.Window, WindowAttr] = weakref.WeakKeyDictionary()
-    """Per-window attributes. I.e., per-session attributes."""
-
     @classmethod
     def required_node_version(cls) -> str:
         """
         Testing playground at https://semver.npmjs.com
         And `0.0.0` means "no restrictions".
         """
-        return ">=14"
+        return ">=14"  # @todo really?
 
     @classmethod
     def should_ignore(cls, view: sublime.View) -> bool:
@@ -81,8 +66,6 @@ def can_start(
     ) -> str | None:
         if message := super().can_start(window, initiating_view, workspace_folders, configuration):
             return message
-
-        cls.window_attrs.setdefault(window, WindowAttr())
         return None
 
     def on_settings_changed(self, settings: DottedDict) -> None:
@@ -154,6 +137,35 @@ def on_server_response_async(self, method: str, response: Response) -> None:
                         documentation["value"] = self.patch_markdown_content(documentation["value"])
             return
 
+    def on_workspace_configuration(self, params: Any, configuration: dict[str, Any]) -> dict[str, Any]:
+        # provide detected venv information from the workspace folder
+        # note that `pyrightconfig.json` seems to be auto-prioritized by the server
+        if (
+            (session := self.weaksession())
+            and (params["section"] == "python")
+            and (scope_uri := params.get("scopeUri"))
+            and (file_path := uri_to_file_path(scope_uri))
+            and (wf_path := find_workspace_folder(session.window, file_path))
+            and (venv_strategies := session.config.settings.get("venvStrategies"))
+            and (venv_info := find_venv_by_finder_names(venv_strategies, project_dir=wf_path))
+        ):
+            self.wf_attrs[wf_path].venv_info = venv_info
+            # When ST just starts, server session hasn't been created yet.
+            # So `on_activated` can't add full information for the initial view and hence we handle it here.
+            if active_view := sublime.active_window().active_view():
+                update_view_status_bar_text(self.__class__, active_view)
+
+            # modify configuration for the venv
+            site_packages_dir = str(venv_info.site_packages_dir)
+            conf_analysis: dict[str, Any] = configuration.setdefault("analysis", {})
+            conf_analysis_extra_paths: list[str] = conf_analysis.setdefault("extraPaths", [])
+            if site_packages_dir not in conf_analysis_extra_paths:
+                conf_analysis_extra_paths.insert(0, site_packages_dir)
+            if not configuration.get("pythonPath"):
+                configuration["pythonPath"] = str(venv_info.python_executable)
+
+        return configuration
+
     # -------------- #
     # custom methods #
     # -------------- #
@@ -171,32 +183,6 @@ def copy_overwrite_dirs(cls) -> None:
         except OSError:
             raise RuntimeError(f'Failed to copy overwrite dirs from "{dir_src}" to "{dir_dst}".')
 
-    def update_status_bar_text(self, extra_variables: dict[str, Any] | None = None) -> None:
-        if not (session := self.weaksession()):
-            return
-
-        variables: dict[str, Any] = {
-            "server_version": self.server_version,
-        }
-
-        if venv_info := self.window_attrs[session.window].venv_info:
-            variables["venv"] = {
-                "finder_name": venv_info.meta.finder_name,
-                "python_version": venv_info.python_version,
-                "venv_prompt": venv_info.prompt,
-            }
-
-        if extra_variables:
-            variables.update(extra_variables)
-
-        rendered_text = ""
-        if template_text := str(session.config.settings.get("statusText") or ""):
-            try:
-                rendered_text = load_string_template(template_text).render(variables)
-            except Exception as e:
-                log_warning(f'Invalid "statusText" template: {e}')
-        session.set_config_status_async(rendered_text)
-
     def patch_markdown_content(self, content: str) -> str:
         # add another linebreak before horizontal rule following fenced code block
         content = re.sub("```\n---", "```\n\n---", content)
@@ -218,40 +204,3 @@ def patch_markdown_content(self, content: str) -> str:
     def parse_server_version(cls) -> str:
         lock_file_content = sublime.load_resource(f"Packages/{PACKAGE_NAME}/language-server/package-lock.json")
         return jmespath.search("dependencies.pyright.version", json.loads(lock_file_content)) or ""
-
-    @classmethod
-    def update_venv_info(
-        cls,
-        settings: DottedDict,
-        workspace_folders: list[WorkspaceFolder],
-        *,
-        window: sublime.Window,
-    ) -> None:
-        window_attr = cls.window_attrs[window]
-
-        def _update_simple_python_path() -> None:
-            window_attr.simple_python_executable = None
-
-            if python_path := first_true(("py", "python3", "python"), pred=shutil.which):
-                window_attr.simple_python_executable = Path(python_path)
-
-        def _update_venv_info() -> None:
-            window_attr.venv_info = None
-
-            if python_path := settings.get("python.pythonPath"):
-                window_attr.venv_info = find_venv_by_python_executable(python_path)
-                return
-
-            supported_finder_names = tuple(get_finder_name_mapping().keys())
-            finder_names: list[str] = settings.get("venvStrategies")
-            if invalid_finder_names := sorted(set(finder_names) - set(supported_finder_names)):
-                log_warning(f"The following finder names are not supported: {', '.join(invalid_finder_names)}")
-
-            if workspace_folders and (first_folder := Path(workspace_folders[0].path).resolve()):
-                for folder in (first_folder, *first_folder.parents):
-                    if venv_info := find_venv_by_finder_names(finder_names, project_dir=folder):
-                        window_attr.venv_info = venv_info
-                        return
-
-        _update_simple_python_path()
-        _update_venv_info()

plugin/dev_environment/helpers.py

@@ -5,7 +5,6 @@
 
 from more_itertools import first_true
 
-from ..virtual_env.venv_info import BaseVenvInfo
 from .impl import (
     BlenderDevEnvironmentHandler,
     GdbDevEnvironmentHandler,
@@ -28,14 +27,9 @@ def get_dev_environment_handler(
     *,
     server_dir: str | Path,
     workspace_folders: Sequence[str],
-    venv_info: BaseVenvInfo | None = None,
 ) -> BaseDevEnvironmentHandler | None:
     if handler_cls := find_dev_environment_handler_class(dev_environment):
-        return handler_cls(
-            server_dir=server_dir,
-            workspace_folders=workspace_folders,
-            venv_info=venv_info,
-        )
+        return handler_cls(server_dir=server_dir, workspace_folders=workspace_folders)
     return None
 
 

plugin/dev_environment/impl/sublime_text.py

@@ -19,7 +19,7 @@ def python_version(self) -> tuple[int, int]:
         return (3, 3)
 
     def handle_(self, *, settings: DottedDict) -> None:
-        self._inject_extra_paths(settings=settings, paths=self.find_package_dependency_dirs(), operation="replace")
+        self._inject_extra_paths(settings=settings, paths=self.find_package_dependency_dirs())
 
     def find_package_dependency_dirs(self) -> list[str]:
         dep_dirs = sys.path.copy()

plugin/dev_environment/interfaces.py

@@ -5,27 +5,19 @@
 from typing import Any, Iterable, Literal, Sequence, final
 
 from LSP.plugin.core.collections import DottedDict
+from more_itertools import unique_everseen
 
 from ..constants import SERVER_SETTING_ANALYSIS_EXTRAPATHS, SERVER_SETTING_DEV_ENVIRONMENT
-from ..log import log_info
+from ..log import log_debug
 from ..utils import camel_to_snake, remove_suffix
-from ..virtual_env.venv_info import BaseVenvInfo
 
 
 class BaseDevEnvironmentHandler(ABC):
-    def __init__(
-        self,
-        *,
-        server_dir: str | Path,
-        workspace_folders: Sequence[str],
-        venv_info: BaseVenvInfo | None = None,
-    ) -> None:
+    def __init__(self, *, server_dir: str | Path, workspace_folders: Sequence[str]) -> None:
         self.server_dir = Path(server_dir)
         """The language server directory."""
         self.workspace_folders = workspace_folders
         """The workspace folders."""
-        self.venv_info = venv_info
-        """The virtual environment information."""
 
     @classmethod
     def name(cls) -> str:
@@ -48,9 +40,6 @@ def handle(self, *, settings: DottedDict) -> None:
         """Handle this environment."""
         self.handle_(settings=settings)
 
-        if self.venv_info:
-            self._inject_extra_paths(settings=settings, paths=(self.venv_info.site_packages_dir,))
-
     @abstractmethod
     def handle_(self, *, settings: DottedDict) -> None:
         """Handle this environment. (subclass)"""
@@ -73,5 +62,7 @@ def _inject_extra_paths(
             next_paths = extra_paths
         else:
             raise ValueError(f"Invalid operation: {operation}")
-        log_info(f"Modified extra analysis paths ({operation = }): {paths}")
+
+        next_paths = list(unique_everseen(next_paths, key=Path))  # deduplication
+        log_debug(f'Due to "dev_environment", new "analysis.extraPaths" is ({operation = }): {next_paths}')
         settings.set(SERVER_SETTING_ANALYSIS_EXTRAPATHS, next_paths)

plugin/utils.py

@@ -60,6 +60,13 @@ def get_default_startupinfo() -> Any:
     return None
 
 
+def path_to_resolved_posix(path: str | Path) -> str | None:
+    try:
+        return Path(path).resolve().as_posix()
+    except Exception:
+        return None
+
+
 def run_shell_command(
     command: str | Sequence[str],
     *,

plugin/utils_lsp.py

@@ -0,0 +1,100 @@
+"""Utility functions related to LSP."""
+
+from __future__ import annotations
+
+from abc import ABC
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import sublime
+from LSP.plugin import parse_uri
+from LSP.plugin.core.registry import windows as lsp_windows_registry
+
+from .log import log_warning
+from .template import load_string_template
+from .utils import drop_falsy, path_to_resolved_posix
+from .virtual_env.venv_finder import BaseVenvInfo
+from .vs_marketplace_lsp_utils import VsMarketplaceClientHandler
+
+
+@dataclass
+class WorkspaceFolderAttr:
+    venv_info: BaseVenvInfo | None = None
+    """The information of the virtual environment."""
+
+
+class AbstractLspPythonPlugin(VsMarketplaceClientHandler, ABC):
+    server_version: str
+    """The version of the language server."""
+    wf_attrs: defaultdict[Path, WorkspaceFolderAttr] = defaultdict(WorkspaceFolderAttr)
+    """Per workspace folder attributes."""
+
+
+def find_workspace_folder(window: sublime.Window, path: str | Path) -> Path | None:
+    """Find a workspace folder for the path. The deepest folder wins if there are multiple matches."""
+    if path_ := path_to_resolved_posix(path):
+        for folder in sorted(drop_falsy(map(path_to_resolved_posix, window.folders())), key=len, reverse=True):
+            if f"{path_}/".startswith(f"{folder}/"):
+                return Path(folder)
+    return None
+
+
+def lowercase_drive_letter(path: str) -> str:
+    """Converts the drive letter in the path to lowercase."""
+    if len(path) > 1 and path[1] == ":":
+        return path[0].lower() + path[1:]
+    return path
+
+
+def uri_to_file_path(uri: str) -> str | None:
+    """Converts the URI to its file path if it's of the "file" scheme. Otherwise, `None`."""
+    scheme, path = parse_uri(uri)
+    return path if scheme == "file" else None
+
+
+def update_view_status_bar_text(
+    lsp_cls: type[AbstractLspPythonPlugin],
+    view: sublime.View,
+    *,
+    extra_variables: dict[str, Any] | None = None,
+) -> None:
+    if not (
+        (file_path := view.file_name())
+        and (window := view.window())
+        and (lsp_window_manager := lsp_windows_registry.lookup(window))
+        and (session := lsp_window_manager.get_session(lsp_cls.name(), file_path))
+    ):
+        return
+
+    # shortcut if the user doesn't want any status text
+    if not (template_text := str(session.config.settings.get("statusText") or "")):
+        session.set_config_status_async("")
+        return
+
+    variables: dict[str, Any] = {
+        "server_version": lsp_cls.server_version,
+    }
+
+    if (
+        (wf_path := find_workspace_folder(window, file_path))
+        and (wf_attr := lsp_cls.wf_attrs.get(wf_path))
+        and (venv_info := wf_attr.venv_info)
+    ):
+        variables["venv"] = {
+            "finder_name": venv_info.meta.finder_name,
+            "python_version": venv_info.python_version,
+            "venv_prompt": venv_info.prompt,
+        }
+
+    if extra_variables:
+        variables.update(extra_variables)
+
+    rendered_text = ""
+    try:
+        rendered_text = load_string_template(template_text).render(variables)
+    except Exception as e:
+        log_warning(f'Invalid "statusText" template: {e}')
+
+    session.set_config_status_async(rendered_text)