diff --git a/LSP-pyright.sublime-settings b/LSP-pyright.sublime-settings index cc434b0..bfbcfcb 100644 --- a/LSP-pyright.sublime-settings +++ b/LSP-pyright.sublime-settings @@ -17,6 +17,18 @@ // - "sublime_text_33": Similar to "sublime_text" but Python 3.3 forced. // - "sublime_text_38": Similar to "sublime_text" but Python 3.8 forced. "pyright.dev_environment": "", + // The strategies used to find a virtual environment in order. + "pyright.venv.strategies": [ + "local_dot_venv", + "env_var_conda_prefix", + "env_var_virtual_env", + "rye", + "poetry", + "pdm", + "pipenv", + "pyenv", + "any_subdirectory", + ], // Offer auto-import completions. "python.analysis.autoImportCompletions": true, // Automatically add common search paths like 'src'? diff --git a/constants.py b/constants.py new file mode 100644 index 0000000..ce36a82 --- /dev/null +++ b/constants.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +assert __package__ + +PACKAGE_NAME = __package__.partition(".")[0] diff --git a/dependencies.json b/dependencies.json index 3d55a18..9526366 100644 --- a/dependencies.json +++ b/dependencies.json @@ -2,7 +2,9 @@ "*": { "*": [ "lsp_utils", - "sublime_lib" + "more-itertools", + "sublime_lib", + "typing-extensions" ] } } diff --git a/log.py b/log.py new file mode 100644 index 0000000..125dd3c --- /dev/null +++ b/log.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from .constants import PACKAGE_NAME + + +def log_debug(message: str) -> None: + print(f"[{PACKAGE_NAME}][DEBUG] {message}") + + +def log_info(message: str) -> None: + print(f"[{PACKAGE_NAME}][INFO] {message}") + + +def log_warning(message: str) -> None: + print(f"[{PACKAGE_NAME}][WARNING] {message}") + + +def log_error(message: str) -> None: + print(f"[{PACKAGE_NAME}][ERROR] {message}") diff --git a/plugin.py b/plugin.py index 54672e6..7fa93c6 100644 --- a/plugin.py +++ b/plugin.py @@ -5,17 +5,21 @@ import shutil import subprocess import sys -from collections.abc import Callable +from collections import defaultdict +from dataclasses import dataclass from pathlib import Path -from typing import Any, cast +from typing import Any, Sequence, cast 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 sublime_lib import ResourcePath -assert __package__ +from .constants import PACKAGE_NAME +from .log import log_info, log_warning +from .venv_finder import VenvInfo, find_venv_by_finder_names, get_finder_name_mapping def plugin_loaded() -> None: @@ -36,11 +40,19 @@ def get_default_startupinfo() -> Any: return None +@dataclass +class WindowAttr: + venv_info: VenvInfo | None = None + + class LspPyrightPlugin(NpmClientHandler): - package_name = __package__.partition(".")[0] + package_name = PACKAGE_NAME server_directory = "language-server" server_binary_path = os.path.join(server_directory, "node_modules", "pyright", "langserver.index.js") + window_attrs: defaultdict[int, WindowAttr] = defaultdict(WindowAttr) + """Per-window attributes. I.e., per-session attributes. The key is the window ID.""" + @classmethod def required_node_version(cls) -> str: """ @@ -62,6 +74,8 @@ def on_settings_changed(self, settings: DottedDict) -> None: settings.set("python.analysis.extraPaths", extraPaths) + self.update_status_bar_text() + @classmethod def on_pre_start( cls, @@ -72,9 +86,10 @@ def on_pre_start( ) -> str | None: super().on_pre_start(window, initiating_view, workspace_folders, configuration) - python_path = cls.python_path(configuration.settings, workspace_folders) - print(f'{cls.name()}: INFO: Using python path "{python_path}"') - configuration.settings.set("python.pythonPath", python_path) + cls.update_venv_info(configuration.settings, workspace_folders, window=window) + if venv_info := cls.window_attrs[window.id()].venv_info: + log_info(f"Using python executable: {venv_info.python_executable}") + configuration.settings.set("python.pythonPath", str(venv_info.python_executable)) return None @classmethod @@ -118,6 +133,24 @@ def on_server_response_async(self, method: str, response: Response) -> None: # custom methods # # -------------- # + def update_status_bar_text(self) -> None: + status_parts: list[str] = [] + + if not (session := self.weaksession()): + return + window_id = session.window.id() + + if venv_info := self.window_attrs[window_id].venv_info: + if venv_info.prompt: + status_parts.append(f"venv: {venv_info.prompt}") + if venv_info.python_version: + status_parts.append(f"py: {venv_info.python_version}") + if venv_info.venv_founder: + status_parts.append(f"by: {venv_info.venv_founder}") + + if session := self.weaksession(): + session.set_config_status_async("; ".join(status_parts)) + 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) @@ -180,77 +213,40 @@ def find_package_dependency_dirs(self, py_ver: tuple[int, int] = (3, 3)) -> list return list(filter(os.path.isdir, dep_dirs)) @classmethod - def python_path(cls, settings: DottedDict, workspace_folders: list[WorkspaceFolder]) -> str: + def update_venv_info( + cls, + settings: DottedDict, + workspace_folders: list[WorkspaceFolder], + *, + window: sublime.Window, + ) -> None: + window_id = window.id() + if python_path := settings.get("python.pythonPath"): - return python_path + cls.window_attrs[window_id].venv_info = VenvInfo.from_python_executable(python_path) + return + + supported_finder_names = tuple(get_finder_name_mapping().keys()) + venv_strategies: Sequence[str] | None = settings.get("pyright.venv.strategies") + if venv_strategies is None: + venv_strategies = supported_finder_names + + if invalid_finder_names := set(venv_strategies) - 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 python_path := cls.python_path_from_venv(folder): - return str(python_path) + if venv_info := find_venv_by_finder_names(venv_strategies, project_dir=folder): + cls.window_attrs[window_id].venv_info = venv_info + return - return shutil.which("python") or shutil.which("python3") or "" + if python_path := shutil.which("python") or shutil.which("python3") or "": + cls.window_attrs[window_id].venv_info = VenvInfo.from_python_executable(python_path) + return - @classmethod - def python_path_from_venv(cls, workspace_folder: Path) -> Path | None: - """ - Resolves the python binary path depending on environment variables and files in the workspace. + cls.window_attrs[window_id].venv_info = None - @see https://github.com/fannheyward/coc-pyright/blob/d58a468b1d7479a1b56906e386f44b997181e307/src/configSettings.ts#L47 - """ - def binary_from_python_path(path: str | Path) -> Path | None: - path = Path(path) - if sublime.platform() == "windows": - binary_path = path / "Scripts/python.exe" - else: - binary_path = path / "bin/python" - return binary_path if binary_path.is_file() else None - - # Config file, venv resolution command, post-processing - venv_config_files: list[tuple[str, str, Callable[[str | Path], Path | None] | None]] = [ - (".pdm-python", "pdm info --python", None), - (".python-version", "pyenv which python", None), - ("Pipfile", "pipenv --py", None), - ("poetry.lock", "poetry env info -p", binary_from_python_path), - ] - - for config_file, command, post_processing in venv_config_files: - if not (workspace_folder / config_file).is_file(): - continue - print(f"{cls.name()}: INFO: {config_file} detected. Run subprocess command: {command}") - try: - stdout, stderr = map( - str.rstrip, - subprocess.Popen( - command, - cwd=workspace_folder, - shell=True, - startupinfo=get_default_startupinfo(), - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - universal_newlines=True, - ).communicate(), - ) - if stderr: - print(f"{cls.name()}: INFO: subprocess stderr: {stderr}") - python_path = stdout - if post_processing: - python_path = post_processing(python_path) - if python_path: - return Path(python_path) - except FileNotFoundError: - print(f"{cls.name()}: WARN: subprocess failed with file not found: {command[0]}") - except PermissionError as e: - print(f"{cls.name()}: WARN: subprocess failed with permission error: {e}") - except subprocess.CalledProcessError as e: - print(f"{cls.name()}: WARN: subprocess failed: {str(e.output).strip()}") - - # virtual environment as subfolder in project - for maybe_venv_path in workspace_folder.iterdir(): - try: - if (maybe_venv_path / "pyvenv.cfg").is_file() and (binary := binary_from_python_path(maybe_venv_path)): - return binary # found a venv - except PermissionError: - pass - return None +class PylanceEventListener(sublime_plugin.EventListener): + def on_pre_close_window(self, window: sublime.Window) -> None: + LspPyrightPlugin.window_attrs.pop(window.id(), None) diff --git a/sublime-package.json b/sublime-package.json index 54b7b6f..1fc2144 100644 --- a/sublime-package.json +++ b/sublime-package.json @@ -37,6 +37,35 @@ "Similar to \"sublime_text\" but Python 3.8 forced." ] }, + "pyright.venv.strategies": { + "default": [], + "description": "The strategies used to find a virtual environment in order.", + "items": { + "enum": [ + "any_subdirectory", + "env_var_conda_prefix", + "env_var_virtual_env", + "local_dot_venv", + "pdm", + "pipenv", + "poetry", + "pyenv", + "rye" + ], + "markdownEnumDescriptions": [ + "Finds the virtual environment with any subdirectory.", + "Finds the virtual environment using the `CONDA_PREFIX` environment variable.", + "Finds the virtual environment using the `VIRTUAL_ENV` environment variable.", + "Finds the virtual environment `.venv` or `venv` directory.", + "Finds the virtual environment using `pdm`.", + "Finds the virtual environment using `pipenv`.", + "Finds the virtual environment using `poetry`.", + "Finds the virtual environment using `pyenv`.", + "Finds the virtual environment using `rye`." + ] + }, + "type": "array", + }, "pyright.disableLanguageServices": { "default": false, "description": "Disables type completion, definitions, and references.", diff --git a/venv_finder.py b/venv_finder.py new file mode 100644 index 0000000..8bf9c78 --- /dev/null +++ b/venv_finder.py @@ -0,0 +1,388 @@ +from __future__ import annotations + +import os +import re +import shutil +import subprocess +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from functools import lru_cache +from pathlib import Path +from types import MappingProxyType +from typing import Any, Generator, Iterable, Mapping, Sequence, final + +from more_itertools import first_true +from typing_extensions import Self + +from .log import log_error + + +def find_venv_by_finder_names(finder_names: Sequence[str], *, project_dir: Path) -> VenvInfo | None: + if isinstance(finder_names, str): + finder_names = (finder_names,) + + for finder_name in finder_names: + if ( + (finder_cls := find_finder_class_by_name(finder_name)) + and finder_cls.can_support(project_dir) + and (venv_info := finder_cls(project_dir).find_venv()) + ): + return venv_info + return None + + +@lru_cache +def find_finder_class_by_name(name: str) -> type[BaseVenvFinder] | None: + """Finds the virtual environment finder class by its name.""" + return get_finder_name_mapping().get(name) + + +@lru_cache +def get_finder_name_mapping() -> Mapping[str, type[BaseVenvFinder]]: + """Returns a mapping of virtual environment finder names to their classes.""" + return MappingProxyType({finder_cls.name(): finder_cls for finder_cls in list_venv_finder_classes()}) + + +def list_venv_finder_classes() -> Generator[type[BaseVenvFinder], None, None]: + """Lists all virtual environment finder classes. The order matters.""" + yield LocalDotVenvVenvFinder + yield EnvVarCondaPrefixVenvFinder + yield EnvVarVirtualEnvVenvFinder + yield RyeVenvFinder + yield PoetryVenvFinder + yield PdmVenvFinder + yield PipenvVenvFinder + yield PyenvVenvFinder + yield AnySubdirectoryVenvFinder + + +@dataclass +class VenvInfo: + """The information of the virtual environment.""" + + venv_dir: Path + """The path of the virtual environment directory.""" + + # ----- # + # cache # + # ----- # + + pyvenv_cfg: dict[str, Any] = field(default_factory=dict) + """The parsed results of the `pyvenv.cfg` file.""" + + # -------- # + # metadata # + # -------- # + + venv_founder: str = "" + """The name of the virtual environment founder.""" + + @property + def prompt(self) -> str: + """The prompt of the virtual environment.""" + if prompt := self.pyvenv_cfg.get("prompt", ""): + return prompt + return self.venv_dir.name + + @property + def python_executable(self) -> Path: + """The path of the Python executable of the virtual environment.""" + if os.name == "nt": + return self.venv_dir / "Scripts/python.exe" + return self.venv_dir / "bin/python" + + @property + def python_version(self) -> str: + """The Python version of the virtual environment.""" + # "venv" module uses "version" + if version := self.pyvenv_cfg.get("version", ""): + return str(version) + # "uv" utility uses "version_info" + if version := self.pyvenv_cfg.get("version_info", ""): + return str(version) + return "" + + @property + def pyvenv_cfg_path(self) -> Path: + """The path of the `pyvenv.cfg` file of the virtual environment.""" + return self.venv_dir / "pyvenv.cfg" + + def is_valid(self) -> bool: + """Checks if this virtual environment is valid.""" + return self.venv_dir.is_dir() and self.pyvenv_cfg_path.is_file() and self.python_executable.is_file() + + def refresh_cache(self) -> None: + """Refreshes cached property values.""" + self.pyvenv_cfg = self.parse_pyvenv_cfg(self.pyvenv_cfg_path) + + @classmethod + def from_venv_dir(cls, venv_dir: str | Path) -> Self | None: + try: + venv_dir = Path(venv_dir).resolve() + except Exception: + return None + + if (venv_info := cls(venv_dir=venv_dir)).is_valid(): + venv_info.refresh_cache() + return venv_info + return None + + @classmethod + def from_python_executable(cls, python_executable: str | Path) -> Self | None: + return cls.from_venv_dir(Path(python_executable).parents[1]) + + @staticmethod + def parse_pyvenv_cfg(pyvenv_cfg: Path) -> dict[str, Any]: + def _pythonize_value(val: str) -> Any: + if val.lower() == "true": + return True + if val.lower() == "false": + return False + if val.isdigit(): + return int(val) + try: + return float(val) + except ValueError: + pass + return val + + cfg: dict[str, Any] = {} + if not pyvenv_cfg.is_file(): + return cfg + with pyvenv_cfg.open(encoding="utf-8") as f: + for line in f: + key, sep, value = line.partition("=") + if not sep: + continue + cfg[key.strip()] = _pythonize_value(value.strip()) + return cfg + + +class BaseVenvFinder(ABC): + REQUIRED_BINS: tuple[str, ...] = tuple() + """The required binaries to run the finder. E.g., `("poetry",)`.""" + + def __init__(self, project_dir: Path) -> None: + self.project_dir = project_dir + """The project root directory.""" + + @final + @classmethod + def name(cls) -> str: + name = cls.__name__ + # remove trailing "VenvFinder" + if name.endswith("VenvFinder"): + name = name[: -len("VenvFinder")] + # CamelCase to snake_case + return "".join(f"_{c.lower()}" if c.isupper() else c for c in name).lstrip("_") + + @final + @classmethod + def check_required_bins(cls) -> bool: + return all(shutil.which(bin_name) for bin_name in cls.REQUIRED_BINS) + + @final + @classmethod + def can_support(cls, project_dir: Path) -> bool: + """Check if this class support the given `project_dir`. Public method.""" + return cls.check_required_bins() and cls._can_support(project_dir) + + @final + def find_venv(self) -> VenvInfo | None: + """Find the virtual environment. Public method.""" + if venv_info := self._find_venv(): + venv_info.venv_founder = self.name() + return venv_info + return None + + @classmethod + @abstractmethod + def _can_support(cls, project_dir: Path) -> bool: + """Check if this class support the given `project_dir`. Implement this method by the subclass.""" + + @abstractmethod + def _find_venv(self) -> VenvInfo | None: + """Find the virtual environment. Implement this method by the subclass.""" + + @staticmethod + def _find_from_venv_dir_candidates(candidates: Iterable[Path]) -> VenvInfo | None: + return first_true(map(VenvInfo.from_venv_dir, filter(Path.is_dir, candidates))) + + @staticmethod + def _run_shell_command(command: str, *, cwd: Path | None = None) -> tuple[str, str] | None: + if os.name == "nt": + # do not create a window for the process + startupinfo = subprocess.STARTUPINFO() # type: ignore + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW # type: ignore + startupinfo.wShowWindow = subprocess.SW_HIDE # type: ignore + else: + startupinfo = None + + try: + stdout, stderr = map( + str.rstrip, + subprocess.Popen( + command, + cwd=cwd, + shell=True, + startupinfo=startupinfo, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + ).communicate(), + ) + except Exception as e: + log_error(f"Failed running command ({command}): {e}") + return None + + if stderr: + log_error(f"Failed running command ({command}): {stderr}") + + return stdout, stderr + + +class AnySubdirectoryVenvFinder(BaseVenvFinder): + """Finds the virtual environment with any subdirectory.""" + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return True + + def _find_venv(self) -> VenvInfo | None: + return self._find_from_venv_dir_candidates(self.project_dir.iterdir()) + + +class EnvVarCondaPrefixVenvFinder(BaseVenvFinder): + """Finds the virtual environment using the `CONDA_PREFIX` environment variable.""" + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return True + + def _find_venv(self) -> VenvInfo | None: + if conda_prefix := os.environ.get("CONDA_PREFIX", ""): + return VenvInfo.from_venv_dir(conda_prefix) + return None + + +class EnvVarVirtualEnvVenvFinder(BaseVenvFinder): + """Finds the virtual environment using the `VIRTUAL_ENV` environment variable.""" + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return True + + def _find_venv(self) -> VenvInfo | None: + if virtual_env := os.environ.get("VIRTUAL_ENV", ""): + return VenvInfo.from_venv_dir(virtual_env) + return None + + +class LocalDotVenvVenvFinder(BaseVenvFinder): + """Finds the virtual environment `.venv` or `venv` directory.""" + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return True + + def _find_venv(self) -> VenvInfo | None: + return self._find_from_venv_dir_candidates(( + self.project_dir / ".venv", + self.project_dir / "venv", + )) + + +class PdmVenvFinder(BaseVenvFinder): + """Finds the virtual environment using `pdm`.""" + + REQUIRED_BINS = ("pdm",) + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return (project_dir / ".pdm-python").is_file() + + def _find_venv(self) -> VenvInfo | None: + if not (output := self._run_shell_command("pdm info --python", cwd=self.project_dir)): + return None + python_executable, _ = output + + if not python_executable: + return None + return VenvInfo.from_python_executable(python_executable) + + +class PipenvVenvFinder(BaseVenvFinder): + """Finds the virtual environment using `pipenv`.""" + + REQUIRED_BINS = ("pipenv",) + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return (project_dir / "Pipfile").is_file() + + def _find_venv(self) -> VenvInfo | None: + if not (output := self._run_shell_command("pipenv --py", cwd=self.project_dir)): + return None + python_executable, _ = output + + if not python_executable: + return None + return VenvInfo.from_python_executable(python_executable) + + +class PoetryVenvFinder(BaseVenvFinder): + """Finds the virtual environment using `poetry`.""" + + REQUIRED_BINS = ("poetry",) + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return (project_dir / "poetry.lock").is_file() + + def _find_venv(self) -> VenvInfo | None: + if not (output := self._run_shell_command("poetry env info -p", cwd=self.project_dir)): + return None + venv_dir, _ = output + + if not venv_dir: + return None + return VenvInfo.from_venv_dir(venv_dir) + + +class PyenvVenvFinder(BaseVenvFinder): + """Finds the virtual environment using `pyenv`.""" + + REQUIRED_BINS = ("pyenv",) + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return (project_dir / ".python-version").is_file() + + def _find_venv(self) -> VenvInfo | None: + if not (output := self._run_shell_command("pyenv which python", cwd=self.project_dir)): + return None + python_executable, _ = output + + if not python_executable: + return None + return VenvInfo.from_python_executable(python_executable) + + +class RyeVenvFinder(BaseVenvFinder): + """Finds the virtual environment using `rye`.""" + + REQUIRED_BINS = ("rye",) + + @classmethod + def _can_support(cls, project_dir: Path) -> bool: + return (project_dir / "pyproject.toml").is_file() + + def _find_venv(self) -> VenvInfo | None: + if not (output := self._run_shell_command("rye show", cwd=self.project_dir)): + return None + stdout, _ = output + + if m := re.search(r"^venv: (.*)$", stdout, re.MULTILINE): + venv_dir = m.group(1) + return VenvInfo.from_venv_dir(venv_dir) + return None