python.rs

   1use anyhow::{Context as _, ensure};
   2use anyhow::{Result, anyhow};
   3use async_trait::async_trait;
   4use collections::HashMap;
   5use futures::{AsyncBufReadExt, StreamExt as _};
   6use gpui::{App, AsyncApp, SharedString, Task};
   7use http_client::github::{AssetKind, GitHubLspBinaryVersion, latest_github_release};
   8use language::language_settings::language_settings;
   9use language::{ContextLocation, LanguageToolchainStore, LspInstaller};
  10use language::{ContextProvider, LspAdapter, LspAdapterDelegate};
  11use language::{LanguageName, ManifestName, ManifestProvider, ManifestQuery};
  12use language::{Toolchain, ToolchainList, ToolchainLister, ToolchainMetadata};
  13use lsp::LanguageServerBinary;
  14use lsp::LanguageServerName;
  15use node_runtime::{NodeRuntime, VersionStrategy};
  16use pet_core::Configuration;
  17use pet_core::os_environment::Environment;
  18use pet_core::python_environment::{PythonEnvironment, PythonEnvironmentKind};
  19use pet_virtualenv::is_virtualenv_dir;
  20use project::Fs;
  21use project::lsp_store::language_server_settings;
  22use serde::{Deserialize, Serialize};
  23use serde_json::{Value, json};
  24use smol::lock::OnceCell;
  25use std::cmp::Ordering;
  26use std::env::consts;
  27use util::command::new_smol_command;
  28use util::fs::{make_file_executable, remove_matching};
  29use util::rel_path::RelPath;
  30
  31use http_client::github_download::{GithubBinaryMetadata, download_server_binary};
  32use parking_lot::Mutex;
  33use std::str::FromStr;
  34use std::{
  35    borrow::Cow,
  36    fmt::Write,
  37    path::{Path, PathBuf},
  38    sync::Arc,
  39};
  40use task::{ShellKind, TaskTemplate, TaskTemplates, VariableName};
  41use util::{ResultExt, maybe};
  42
  43#[derive(Debug, Serialize, Deserialize)]
  44pub(crate) struct PythonToolchainData {
  45    #[serde(flatten)]
  46    environment: PythonEnvironment,
  47    #[serde(skip_serializing_if = "Option::is_none")]
  48    activation_scripts: Option<HashMap<ShellKind, PathBuf>>,
  49}
  50
  51pub(crate) struct PyprojectTomlManifestProvider;
  52
  53impl ManifestProvider for PyprojectTomlManifestProvider {
  54    fn name(&self) -> ManifestName {
  55        SharedString::new_static("pyproject.toml").into()
  56    }
  57
  58    fn search(
  59        &self,
  60        ManifestQuery {
  61            path,
  62            depth,
  63            delegate,
  64        }: ManifestQuery,
  65    ) -> Option<Arc<RelPath>> {
  66        for path in path.ancestors().take(depth) {
  67            let p = path.join(RelPath::unix("pyproject.toml").unwrap());
  68            if delegate.exists(&p, Some(false)) {
  69                return Some(path.into());
  70            }
  71        }
  72
  73        None
  74    }
  75}
  76
  77enum TestRunner {
  78    UNITTEST,
  79    PYTEST,
  80}
  81
  82impl FromStr for TestRunner {
  83    type Err = ();
  84
  85    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
  86        match s {
  87            "unittest" => Ok(Self::UNITTEST),
  88            "pytest" => Ok(Self::PYTEST),
  89            _ => Err(()),
  90        }
  91    }
  92}
  93
  94/// Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
  95/// Where `XX` is the sorting category, `YYYY` is based on most recent usage,
  96/// and `name` is the symbol name itself.
  97///
  98/// The problem with it is that Pyright adjusts the sort text based on previous resolutions (items for which we've issued `completion/resolve` call have their sortText adjusted),
  99/// which - long story short - makes completion items list non-stable. Pyright probably relies on VSCode's implementation detail.
 100/// see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
 101fn process_pyright_completions(items: &mut [lsp::CompletionItem]) {
 102    for item in items {
 103        item.sort_text.take();
 104    }
 105}
 106
 107pub struct TyLspAdapter {
 108    fs: Arc<dyn Fs>,
 109}
 110
 111#[cfg(target_os = "macos")]
 112impl TyLspAdapter {
 113    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
 114    const ARCH_SERVER_NAME: &str = "apple-darwin";
 115}
 116
 117#[cfg(target_os = "linux")]
 118impl TyLspAdapter {
 119    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
 120    const ARCH_SERVER_NAME: &str = "unknown-linux-gnu";
 121}
 122
 123#[cfg(target_os = "freebsd")]
 124impl TyLspAdapter {
 125    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
 126    const ARCH_SERVER_NAME: &str = "unknown-freebsd";
 127}
 128
 129#[cfg(target_os = "windows")]
 130impl TyLspAdapter {
 131    const GITHUB_ASSET_KIND: AssetKind = AssetKind::Zip;
 132    const ARCH_SERVER_NAME: &str = "pc-windows-msvc";
 133}
 134
 135impl TyLspAdapter {
 136    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("ty");
 137
 138    pub fn new(fs: Arc<dyn Fs>) -> TyLspAdapter {
 139        TyLspAdapter { fs }
 140    }
 141
 142    fn build_asset_name() -> Result<(String, String)> {
 143        let arch = match consts::ARCH {
 144            "x86" => "i686",
 145            _ => consts::ARCH,
 146        };
 147        let os = Self::ARCH_SERVER_NAME;
 148        let suffix = match consts::OS {
 149            "windows" => "zip",
 150            _ => "tar.gz",
 151        };
 152        let asset_name = format!("ty-{arch}-{os}.{suffix}");
 153        let asset_stem = format!("ty-{arch}-{os}");
 154        Ok((asset_stem, asset_name))
 155    }
 156}
 157
 158#[async_trait(?Send)]
 159impl LspAdapter for TyLspAdapter {
 160    fn name(&self) -> LanguageServerName {
 161        Self::SERVER_NAME
 162    }
 163
 164    async fn workspace_configuration(
 165        self: Arc<Self>,
 166        delegate: &Arc<dyn LspAdapterDelegate>,
 167        toolchain: Option<Toolchain>,
 168        cx: &mut AsyncApp,
 169    ) -> Result<Value> {
 170        let mut ret = cx
 171            .update(|cx| {
 172                language_server_settings(delegate.as_ref(), &self.name(), cx)
 173                    .and_then(|s| s.settings.clone())
 174            })?
 175            .unwrap_or_else(|| json!({}));
 176        if let Some(toolchain) = toolchain.and_then(|toolchain| {
 177            serde_json::from_value::<PythonToolchainData>(toolchain.as_json).ok()
 178        }) {
 179            _ = maybe!({
 180                let uri =
 181                    url::Url::from_file_path(toolchain.environment.executable.as_ref()?).ok()?;
 182                let sys_prefix = toolchain.environment.prefix.clone()?;
 183                let environment = json!({
 184                    "executable": {
 185                        "uri": uri,
 186                        "sysPrefix": sys_prefix
 187                    }
 188                });
 189                ret.as_object_mut()?
 190                    .entry("pythonExtension")
 191                    .or_insert_with(|| json!({ "activeEnvironment": environment }));
 192                Some(())
 193            });
 194        }
 195        Ok(json!({"ty": ret}))
 196    }
 197}
 198
 199impl LspInstaller for TyLspAdapter {
 200    type BinaryVersion = GitHubLspBinaryVersion;
 201    async fn fetch_latest_server_version(
 202        &self,
 203        delegate: &dyn LspAdapterDelegate,
 204        _: bool,
 205        _: &mut AsyncApp,
 206    ) -> Result<Self::BinaryVersion> {
 207        let release =
 208            latest_github_release("astral-sh/ty", true, true, delegate.http_client()).await?;
 209        let (_, asset_name) = Self::build_asset_name()?;
 210        let asset = release
 211            .assets
 212            .into_iter()
 213            .find(|asset| asset.name == asset_name)
 214            .with_context(|| format!("no asset found matching `{asset_name:?}`"))?;
 215        Ok(GitHubLspBinaryVersion {
 216            name: release.tag_name,
 217            url: asset.browser_download_url,
 218            digest: asset.digest,
 219        })
 220    }
 221
 222    async fn fetch_server_binary(
 223        &self,
 224        latest_version: Self::BinaryVersion,
 225        container_dir: PathBuf,
 226        delegate: &dyn LspAdapterDelegate,
 227    ) -> Result<LanguageServerBinary> {
 228        let GitHubLspBinaryVersion {
 229            name,
 230            url,
 231            digest: expected_digest,
 232        } = latest_version;
 233        let destination_path = container_dir.join(format!("ty-{name}"));
 234
 235        async_fs::create_dir_all(&destination_path).await?;
 236
 237        let server_path = match Self::GITHUB_ASSET_KIND {
 238            AssetKind::TarGz | AssetKind::Gz => destination_path
 239                .join(Self::build_asset_name()?.0)
 240                .join("ty"),
 241            AssetKind::Zip => destination_path.clone().join("ty.exe"),
 242        };
 243
 244        let binary = LanguageServerBinary {
 245            path: server_path.clone(),
 246            env: None,
 247            arguments: vec!["server".into()],
 248        };
 249
 250        let metadata_path = destination_path.with_extension("metadata");
 251        let metadata = GithubBinaryMetadata::read_from_file(&metadata_path)
 252            .await
 253            .ok();
 254        if let Some(metadata) = metadata {
 255            let validity_check = async || {
 256                delegate
 257                    .try_exec(LanguageServerBinary {
 258                        path: server_path.clone(),
 259                        arguments: vec!["--version".into()],
 260                        env: None,
 261                    })
 262                    .await
 263                    .inspect_err(|err| {
 264                        log::warn!("Unable to run {server_path:?} asset, redownloading: {err}",)
 265                    })
 266            };
 267            if let (Some(actual_digest), Some(expected_digest)) =
 268                (&metadata.digest, &expected_digest)
 269            {
 270                if actual_digest == expected_digest {
 271                    if validity_check().await.is_ok() {
 272                        return Ok(binary);
 273                    }
 274                } else {
 275                    log::info!(
 276                        "SHA-256 mismatch for {destination_path:?} asset, downloading new asset. Expected: {expected_digest}, Got: {actual_digest}"
 277                    );
 278                }
 279            } else if validity_check().await.is_ok() {
 280                return Ok(binary);
 281            }
 282        }
 283
 284        download_server_binary(
 285            &*delegate.http_client(),
 286            &url,
 287            expected_digest.as_deref(),
 288            &destination_path,
 289            Self::GITHUB_ASSET_KIND,
 290        )
 291        .await?;
 292        make_file_executable(&server_path).await?;
 293        remove_matching(&container_dir, |path| path != destination_path).await;
 294        GithubBinaryMetadata::write_to_file(
 295            &GithubBinaryMetadata {
 296                metadata_version: 1,
 297                digest: expected_digest,
 298            },
 299            &metadata_path,
 300        )
 301        .await?;
 302
 303        Ok(LanguageServerBinary {
 304            path: server_path,
 305            env: None,
 306            arguments: vec!["server".into()],
 307        })
 308    }
 309
 310    async fn cached_server_binary(
 311        &self,
 312        container_dir: PathBuf,
 313        _: &dyn LspAdapterDelegate,
 314    ) -> Option<LanguageServerBinary> {
 315        maybe!(async {
 316            let mut last = None;
 317            let mut entries = self.fs.read_dir(&container_dir).await?;
 318            while let Some(entry) = entries.next().await {
 319                let path = entry?;
 320                if path.extension().is_some_and(|ext| ext == "metadata") {
 321                    continue;
 322                }
 323                last = Some(path);
 324            }
 325
 326            let path = last.context("no cached binary")?;
 327            let path = match TyLspAdapter::GITHUB_ASSET_KIND {
 328                AssetKind::TarGz | AssetKind::Gz => {
 329                    path.join(Self::build_asset_name()?.0).join("ty")
 330                }
 331                AssetKind::Zip => path.join("ty.exe"),
 332            };
 333
 334            anyhow::Ok(LanguageServerBinary {
 335                path,
 336                env: None,
 337                arguments: vec!["server".into()],
 338            })
 339        })
 340        .await
 341        .log_err()
 342    }
 343}
 344
 345pub struct PyrightLspAdapter {
 346    node: NodeRuntime,
 347}
 348
 349impl PyrightLspAdapter {
 350    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pyright");
 351    const SERVER_PATH: &str = "node_modules/pyright/langserver.index.js";
 352    const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "pyright/langserver.index.js";
 353
 354    pub fn new(node: NodeRuntime) -> Self {
 355        PyrightLspAdapter { node }
 356    }
 357
 358    async fn get_cached_server_binary(
 359        container_dir: PathBuf,
 360        node: &NodeRuntime,
 361    ) -> Option<LanguageServerBinary> {
 362        let server_path = container_dir.join(Self::SERVER_PATH);
 363        if server_path.exists() {
 364            Some(LanguageServerBinary {
 365                path: node.binary_path().await.log_err()?,
 366                env: None,
 367                arguments: vec![server_path.into(), "--stdio".into()],
 368            })
 369        } else {
 370            log::error!("missing executable in directory {:?}", server_path);
 371            None
 372        }
 373    }
 374}
 375
 376#[async_trait(?Send)]
 377impl LspAdapter for PyrightLspAdapter {
 378    fn name(&self) -> LanguageServerName {
 379        Self::SERVER_NAME
 380    }
 381
 382    async fn initialization_options(
 383        self: Arc<Self>,
 384        _: &Arc<dyn LspAdapterDelegate>,
 385    ) -> Result<Option<Value>> {
 386        // Provide minimal initialization options
 387        // Virtual environment configuration will be handled through workspace configuration
 388        Ok(Some(json!({
 389            "python": {
 390                "analysis": {
 391                    "autoSearchPaths": true,
 392                    "useLibraryCodeForTypes": true,
 393                    "autoImportCompletions": true
 394                }
 395            }
 396        })))
 397    }
 398
 399    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
 400        process_pyright_completions(items);
 401    }
 402
 403    async fn label_for_completion(
 404        &self,
 405        item: &lsp::CompletionItem,
 406        language: &Arc<language::Language>,
 407    ) -> Option<language::CodeLabel> {
 408        let label = &item.label;
 409        let label_len = label.len();
 410        let grammar = language.grammar()?;
 411        let highlight_id = match item.kind? {
 412            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method"),
 413            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function"),
 414            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type"),
 415            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant"),
 416            lsp::CompletionItemKind::VARIABLE => grammar.highlight_id_for_name("variable"),
 417            _ => {
 418                return None;
 419            }
 420        };
 421        let mut text = label.clone();
 422        if let Some(completion_details) = item
 423            .label_details
 424            .as_ref()
 425            .and_then(|details| details.description.as_ref())
 426        {
 427            write!(&mut text, " {}", completion_details).ok();
 428        }
 429        Some(language::CodeLabel::filtered(
 430            text,
 431            label_len,
 432            item.filter_text.as_deref(),
 433            highlight_id
 434                .map(|id| (0..label_len, id))
 435                .into_iter()
 436                .collect(),
 437        ))
 438    }
 439
 440    async fn label_for_symbol(
 441        &self,
 442        name: &str,
 443        kind: lsp::SymbolKind,
 444        language: &Arc<language::Language>,
 445    ) -> Option<language::CodeLabel> {
 446        let (text, filter_range, display_range) = match kind {
 447            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
 448                let text = format!("def {}():\n", name);
 449                let filter_range = 4..4 + name.len();
 450                let display_range = 0..filter_range.end;
 451                (text, filter_range, display_range)
 452            }
 453            lsp::SymbolKind::CLASS => {
 454                let text = format!("class {}:", name);
 455                let filter_range = 6..6 + name.len();
 456                let display_range = 0..filter_range.end;
 457                (text, filter_range, display_range)
 458            }
 459            lsp::SymbolKind::CONSTANT => {
 460                let text = format!("{} = 0", name);
 461                let filter_range = 0..name.len();
 462                let display_range = 0..filter_range.end;
 463                (text, filter_range, display_range)
 464            }
 465            _ => return None,
 466        };
 467
 468        Some(language::CodeLabel::new(
 469            text[display_range.clone()].to_string(),
 470            filter_range,
 471            language.highlight_text(&text.as_str().into(), display_range),
 472        ))
 473    }
 474
 475    async fn workspace_configuration(
 476        self: Arc<Self>,
 477        adapter: &Arc<dyn LspAdapterDelegate>,
 478        toolchain: Option<Toolchain>,
 479        cx: &mut AsyncApp,
 480    ) -> Result<Value> {
 481        cx.update(move |cx| {
 482            let mut user_settings =
 483                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
 484                    .and_then(|s| s.settings.clone())
 485                    .unwrap_or_default();
 486
 487            // If we have a detected toolchain, configure Pyright to use it
 488            if let Some(toolchain) = toolchain
 489                && let Ok(env) =
 490                    serde_json::from_value::<PythonToolchainData>(toolchain.as_json.clone())
 491            {
 492                if !user_settings.is_object() {
 493                    user_settings = Value::Object(serde_json::Map::default());
 494                }
 495                let object = user_settings.as_object_mut().unwrap();
 496
 497                let interpreter_path = toolchain.path.to_string();
 498                if let Some(venv_dir) = &env.environment.prefix {
 499                    // Set venvPath and venv at the root level
 500                    // This matches the format of a pyrightconfig.json file
 501                    if let Some(parent) = venv_dir.parent() {
 502                        // Use relative path if the venv is inside the workspace
 503                        let venv_path = if parent == adapter.worktree_root_path() {
 504                            ".".to_string()
 505                        } else {
 506                            parent.to_string_lossy().into_owned()
 507                        };
 508                        object.insert("venvPath".to_string(), Value::String(venv_path));
 509                    }
 510
 511                    if let Some(venv_name) = venv_dir.file_name() {
 512                        object.insert(
 513                            "venv".to_owned(),
 514                            Value::String(venv_name.to_string_lossy().into_owned()),
 515                        );
 516                    }
 517                }
 518
 519                // Always set the python interpreter path
 520                // Get or create the python section
 521                let python = object
 522                    .entry("python")
 523                    .and_modify(|v| {
 524                        if !v.is_object() {
 525                            *v = Value::Object(serde_json::Map::default());
 526                        }
 527                    })
 528                    .or_insert(Value::Object(serde_json::Map::default()));
 529                let python = python.as_object_mut().unwrap();
 530
 531                // Set both pythonPath and defaultInterpreterPath for compatibility
 532                python.insert(
 533                    "pythonPath".to_owned(),
 534                    Value::String(interpreter_path.clone()),
 535                );
 536                python.insert(
 537                    "defaultInterpreterPath".to_owned(),
 538                    Value::String(interpreter_path),
 539                );
 540            }
 541
 542            user_settings
 543        })
 544    }
 545}
 546
 547impl LspInstaller for PyrightLspAdapter {
 548    type BinaryVersion = String;
 549
 550    async fn fetch_latest_server_version(
 551        &self,
 552        _: &dyn LspAdapterDelegate,
 553        _: bool,
 554        _: &mut AsyncApp,
 555    ) -> Result<String> {
 556        self.node
 557            .npm_package_latest_version(Self::SERVER_NAME.as_ref())
 558            .await
 559    }
 560
 561    async fn check_if_user_installed(
 562        &self,
 563        delegate: &dyn LspAdapterDelegate,
 564        _: Option<Toolchain>,
 565        _: &AsyncApp,
 566    ) -> Option<LanguageServerBinary> {
 567        if let Some(pyright_bin) = delegate.which("pyright-langserver".as_ref()).await {
 568            let env = delegate.shell_env().await;
 569            Some(LanguageServerBinary {
 570                path: pyright_bin,
 571                env: Some(env),
 572                arguments: vec!["--stdio".into()],
 573            })
 574        } else {
 575            let node = delegate.which("node".as_ref()).await?;
 576            let (node_modules_path, _) = delegate
 577                .npm_package_installed_version(Self::SERVER_NAME.as_ref())
 578                .await
 579                .log_err()??;
 580
 581            let path = node_modules_path.join(Self::NODE_MODULE_RELATIVE_SERVER_PATH);
 582
 583            let env = delegate.shell_env().await;
 584            Some(LanguageServerBinary {
 585                path: node,
 586                env: Some(env),
 587                arguments: vec![path.into(), "--stdio".into()],
 588            })
 589        }
 590    }
 591
 592    async fn fetch_server_binary(
 593        &self,
 594        latest_version: Self::BinaryVersion,
 595        container_dir: PathBuf,
 596        delegate: &dyn LspAdapterDelegate,
 597    ) -> Result<LanguageServerBinary> {
 598        let server_path = container_dir.join(Self::SERVER_PATH);
 599
 600        self.node
 601            .npm_install_packages(
 602                &container_dir,
 603                &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
 604            )
 605            .await?;
 606
 607        let env = delegate.shell_env().await;
 608        Ok(LanguageServerBinary {
 609            path: self.node.binary_path().await?,
 610            env: Some(env),
 611            arguments: vec![server_path.into(), "--stdio".into()],
 612        })
 613    }
 614
 615    async fn check_if_version_installed(
 616        &self,
 617        version: &Self::BinaryVersion,
 618        container_dir: &PathBuf,
 619        delegate: &dyn LspAdapterDelegate,
 620    ) -> Option<LanguageServerBinary> {
 621        let server_path = container_dir.join(Self::SERVER_PATH);
 622
 623        let should_install_language_server = self
 624            .node
 625            .should_install_npm_package(
 626                Self::SERVER_NAME.as_ref(),
 627                &server_path,
 628                container_dir,
 629                VersionStrategy::Latest(version),
 630            )
 631            .await;
 632
 633        if should_install_language_server {
 634            None
 635        } else {
 636            let env = delegate.shell_env().await;
 637            Some(LanguageServerBinary {
 638                path: self.node.binary_path().await.ok()?,
 639                env: Some(env),
 640                arguments: vec![server_path.into(), "--stdio".into()],
 641            })
 642        }
 643    }
 644
 645    async fn cached_server_binary(
 646        &self,
 647        container_dir: PathBuf,
 648        delegate: &dyn LspAdapterDelegate,
 649    ) -> Option<LanguageServerBinary> {
 650        let mut binary = Self::get_cached_server_binary(container_dir, &self.node).await?;
 651        binary.env = Some(delegate.shell_env().await);
 652        Some(binary)
 653    }
 654}
 655
 656pub(crate) struct PythonContextProvider;
 657
 658const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
 659    VariableName::Custom(Cow::Borrowed("PYTHON_TEST_TARGET"));
 660
 661const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
 662    VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
 663
 664const PYTHON_MODULE_NAME_TASK_VARIABLE: VariableName =
 665    VariableName::Custom(Cow::Borrowed("PYTHON_MODULE_NAME"));
 666
 667impl ContextProvider for PythonContextProvider {
 668    fn build_context(
 669        &self,
 670        variables: &task::TaskVariables,
 671        location: ContextLocation<'_>,
 672        _: Option<HashMap<String, String>>,
 673        toolchains: Arc<dyn LanguageToolchainStore>,
 674        cx: &mut gpui::App,
 675    ) -> Task<Result<task::TaskVariables>> {
 676        let test_target =
 677            match selected_test_runner(location.file_location.buffer.read(cx).file(), cx) {
 678                TestRunner::UNITTEST => self.build_unittest_target(variables),
 679                TestRunner::PYTEST => self.build_pytest_target(variables),
 680            };
 681
 682        let module_target = self.build_module_target(variables);
 683        let location_file = location.file_location.buffer.read(cx).file().cloned();
 684        let worktree_id = location_file.as_ref().map(|f| f.worktree_id(cx));
 685
 686        cx.spawn(async move |cx| {
 687            let active_toolchain = if let Some(worktree_id) = worktree_id {
 688                let file_path = location_file
 689                    .as_ref()
 690                    .and_then(|f| f.path().parent())
 691                    .map(Arc::from)
 692                    .unwrap_or_else(|| RelPath::empty().into());
 693
 694                toolchains
 695                    .active_toolchain(worktree_id, file_path, "Python".into(), cx)
 696                    .await
 697                    .map_or_else(
 698                        || String::from("python3"),
 699                        |toolchain| toolchain.path.to_string(),
 700                    )
 701            } else {
 702                String::from("python3")
 703            };
 704
 705            let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
 706
 707            Ok(task::TaskVariables::from_iter(
 708                test_target
 709                    .into_iter()
 710                    .chain(module_target.into_iter())
 711                    .chain([toolchain]),
 712            ))
 713        })
 714    }
 715
 716    fn associated_tasks(
 717        &self,
 718        file: Option<Arc<dyn language::File>>,
 719        cx: &App,
 720    ) -> Task<Option<TaskTemplates>> {
 721        let test_runner = selected_test_runner(file.as_ref(), cx);
 722
 723        let mut tasks = vec![
 724            // Execute a selection
 725            TaskTemplate {
 726                label: "execute selection".to_owned(),
 727                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 728                args: vec![
 729                    "-c".to_owned(),
 730                    VariableName::SelectedText.template_value_with_whitespace(),
 731                ],
 732                cwd: Some(VariableName::WorktreeRoot.template_value()),
 733                ..TaskTemplate::default()
 734            },
 735            // Execute an entire file
 736            TaskTemplate {
 737                label: format!("run '{}'", VariableName::File.template_value()),
 738                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 739                args: vec![VariableName::File.template_value_with_whitespace()],
 740                cwd: Some(VariableName::WorktreeRoot.template_value()),
 741                ..TaskTemplate::default()
 742            },
 743            // Execute a file as module
 744            TaskTemplate {
 745                label: format!("run module '{}'", VariableName::File.template_value()),
 746                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 747                args: vec![
 748                    "-m".to_owned(),
 749                    PYTHON_MODULE_NAME_TASK_VARIABLE.template_value(),
 750                ],
 751                cwd: Some(VariableName::WorktreeRoot.template_value()),
 752                tags: vec!["python-module-main-method".to_owned()],
 753                ..TaskTemplate::default()
 754            },
 755        ];
 756
 757        tasks.extend(match test_runner {
 758            TestRunner::UNITTEST => {
 759                [
 760                    // Run tests for an entire file
 761                    TaskTemplate {
 762                        label: format!("unittest '{}'", VariableName::File.template_value()),
 763                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 764                        args: vec![
 765                            "-m".to_owned(),
 766                            "unittest".to_owned(),
 767                            VariableName::File.template_value_with_whitespace(),
 768                        ],
 769                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 770                        ..TaskTemplate::default()
 771                    },
 772                    // Run test(s) for a specific target within a file
 773                    TaskTemplate {
 774                        label: "unittest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 775                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 776                        args: vec![
 777                            "-m".to_owned(),
 778                            "unittest".to_owned(),
 779                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 780                        ],
 781                        tags: vec![
 782                            "python-unittest-class".to_owned(),
 783                            "python-unittest-method".to_owned(),
 784                        ],
 785                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 786                        ..TaskTemplate::default()
 787                    },
 788                ]
 789            }
 790            TestRunner::PYTEST => {
 791                [
 792                    // Run tests for an entire file
 793                    TaskTemplate {
 794                        label: format!("pytest '{}'", VariableName::File.template_value()),
 795                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 796                        args: vec![
 797                            "-m".to_owned(),
 798                            "pytest".to_owned(),
 799                            VariableName::File.template_value_with_whitespace(),
 800                        ],
 801                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 802                        ..TaskTemplate::default()
 803                    },
 804                    // Run test(s) for a specific target within a file
 805                    TaskTemplate {
 806                        label: "pytest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 807                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 808                        args: vec![
 809                            "-m".to_owned(),
 810                            "pytest".to_owned(),
 811                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 812                        ],
 813                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 814                        tags: vec![
 815                            "python-pytest-class".to_owned(),
 816                            "python-pytest-method".to_owned(),
 817                        ],
 818                        ..TaskTemplate::default()
 819                    },
 820                ]
 821            }
 822        });
 823
 824        Task::ready(Some(TaskTemplates(tasks)))
 825    }
 826}
 827
 828fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &App) -> TestRunner {
 829    const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
 830    language_settings(Some(LanguageName::new("Python")), location, cx)
 831        .tasks
 832        .variables
 833        .get(TEST_RUNNER_VARIABLE)
 834        .and_then(|val| TestRunner::from_str(val).ok())
 835        .unwrap_or(TestRunner::PYTEST)
 836}
 837
 838impl PythonContextProvider {
 839    fn build_unittest_target(
 840        &self,
 841        variables: &task::TaskVariables,
 842    ) -> Option<(VariableName, String)> {
 843        let python_module_name =
 844            python_module_name_from_relative_path(variables.get(&VariableName::RelativeFile)?);
 845
 846        let unittest_class_name =
 847            variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
 848
 849        let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
 850            "_unittest_method_name",
 851        )));
 852
 853        let unittest_target_str = match (unittest_class_name, unittest_method_name) {
 854            (Some(class_name), Some(method_name)) => {
 855                format!("{python_module_name}.{class_name}.{method_name}")
 856            }
 857            (Some(class_name), None) => format!("{python_module_name}.{class_name}"),
 858            (None, None) => python_module_name,
 859            // should never happen, a TestCase class is the unit of testing
 860            (None, Some(_)) => return None,
 861        };
 862
 863        Some((
 864            PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
 865            unittest_target_str,
 866        ))
 867    }
 868
 869    fn build_pytest_target(
 870        &self,
 871        variables: &task::TaskVariables,
 872    ) -> Option<(VariableName, String)> {
 873        let file_path = variables.get(&VariableName::RelativeFile)?;
 874
 875        let pytest_class_name =
 876            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
 877
 878        let pytest_method_name =
 879            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
 880
 881        let pytest_target_str = match (pytest_class_name, pytest_method_name) {
 882            (Some(class_name), Some(method_name)) => {
 883                format!("{file_path}::{class_name}::{method_name}")
 884            }
 885            (Some(class_name), None) => {
 886                format!("{file_path}::{class_name}")
 887            }
 888            (None, Some(method_name)) => {
 889                format!("{file_path}::{method_name}")
 890            }
 891            (None, None) => file_path.to_string(),
 892        };
 893
 894        Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
 895    }
 896
 897    fn build_module_target(
 898        &self,
 899        variables: &task::TaskVariables,
 900    ) -> Result<(VariableName, String)> {
 901        let python_module_name = python_module_name_from_relative_path(
 902            variables.get(&VariableName::RelativeFile).unwrap_or(""),
 903        );
 904
 905        let module_target = (PYTHON_MODULE_NAME_TASK_VARIABLE.clone(), python_module_name);
 906
 907        Ok(module_target)
 908    }
 909}
 910
 911fn python_module_name_from_relative_path(relative_path: &str) -> String {
 912    let path_with_dots = relative_path.replace('/', ".");
 913    path_with_dots
 914        .strip_suffix(".py")
 915        .unwrap_or(&path_with_dots)
 916        .to_string()
 917}
 918
 919fn is_python_env_global(k: &PythonEnvironmentKind) -> bool {
 920    matches!(
 921        k,
 922        PythonEnvironmentKind::Homebrew
 923            | PythonEnvironmentKind::Pyenv
 924            | PythonEnvironmentKind::GlobalPaths
 925            | PythonEnvironmentKind::MacPythonOrg
 926            | PythonEnvironmentKind::MacCommandLineTools
 927            | PythonEnvironmentKind::LinuxGlobal
 928            | PythonEnvironmentKind::MacXCode
 929            | PythonEnvironmentKind::WindowsStore
 930            | PythonEnvironmentKind::WindowsRegistry
 931    )
 932}
 933
 934fn python_env_kind_display(k: &PythonEnvironmentKind) -> &'static str {
 935    match k {
 936        PythonEnvironmentKind::Conda => "Conda",
 937        PythonEnvironmentKind::Pixi => "pixi",
 938        PythonEnvironmentKind::Homebrew => "Homebrew",
 939        PythonEnvironmentKind::Pyenv => "global (Pyenv)",
 940        PythonEnvironmentKind::GlobalPaths => "global",
 941        PythonEnvironmentKind::PyenvVirtualEnv => "Pyenv",
 942        PythonEnvironmentKind::Pipenv => "Pipenv",
 943        PythonEnvironmentKind::Poetry => "Poetry",
 944        PythonEnvironmentKind::MacPythonOrg => "global (Python.org)",
 945        PythonEnvironmentKind::MacCommandLineTools => "global (Command Line Tools for Xcode)",
 946        PythonEnvironmentKind::LinuxGlobal => "global",
 947        PythonEnvironmentKind::MacXCode => "global (Xcode)",
 948        PythonEnvironmentKind::Venv => "venv",
 949        PythonEnvironmentKind::VirtualEnv => "virtualenv",
 950        PythonEnvironmentKind::VirtualEnvWrapper => "virtualenvwrapper",
 951        PythonEnvironmentKind::WindowsStore => "global (Windows Store)",
 952        PythonEnvironmentKind::WindowsRegistry => "global (Windows Registry)",
 953    }
 954}
 955
 956pub(crate) struct PythonToolchainProvider;
 957
 958static ENV_PRIORITY_LIST: &[PythonEnvironmentKind] = &[
 959    // Prioritize non-Conda environments.
 960    PythonEnvironmentKind::Poetry,
 961    PythonEnvironmentKind::Pipenv,
 962    PythonEnvironmentKind::VirtualEnvWrapper,
 963    PythonEnvironmentKind::Venv,
 964    PythonEnvironmentKind::VirtualEnv,
 965    PythonEnvironmentKind::PyenvVirtualEnv,
 966    PythonEnvironmentKind::Pixi,
 967    PythonEnvironmentKind::Conda,
 968    PythonEnvironmentKind::Pyenv,
 969    PythonEnvironmentKind::GlobalPaths,
 970    PythonEnvironmentKind::Homebrew,
 971];
 972
 973fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
 974    if let Some(kind) = kind {
 975        ENV_PRIORITY_LIST
 976            .iter()
 977            .position(|blessed_env| blessed_env == &kind)
 978            .unwrap_or(ENV_PRIORITY_LIST.len())
 979    } else {
 980        // Unknown toolchains are less useful than non-blessed ones.
 981        ENV_PRIORITY_LIST.len() + 1
 982    }
 983}
 984
 985/// Return the name of environment declared in <worktree-root/.venv.
 986///
 987/// https://virtualfish.readthedocs.io/en/latest/plugins.html#auto-activation-auto-activation
 988async fn get_worktree_venv_declaration(worktree_root: &Path) -> Option<String> {
 989    let file = async_fs::File::open(worktree_root.join(".venv"))
 990        .await
 991        .ok()?;
 992    let mut venv_name = String::new();
 993    smol::io::BufReader::new(file)
 994        .read_line(&mut venv_name)
 995        .await
 996        .ok()?;
 997    Some(venv_name.trim().to_string())
 998}
 999
1000fn get_venv_parent_dir(env: &PythonEnvironment) -> Option<PathBuf> {
1001    // If global, we aren't a virtual environment
1002    if let Some(kind) = env.kind
1003        && is_python_env_global(&kind)
1004    {
1005        return None;
1006    }
1007
1008    // Check to be sure we are a virtual environment using pet's most generic
1009    // virtual environment type, VirtualEnv
1010    let venv = env
1011        .executable
1012        .as_ref()
1013        .and_then(|p| p.parent())
1014        .and_then(|p| p.parent())
1015        .filter(|p| is_virtualenv_dir(p))?;
1016
1017    venv.parent().map(|parent| parent.to_path_buf())
1018}
1019
1020fn wr_distance(wr: &PathBuf, venv: Option<&PathBuf>) -> usize {
1021    if let Some(venv) = venv
1022        && let Ok(p) = venv.strip_prefix(wr)
1023    {
1024        p.components().count()
1025    } else {
1026        usize::MAX
1027    }
1028}
1029
1030#[async_trait]
1031impl ToolchainLister for PythonToolchainProvider {
1032    async fn list(
1033        &self,
1034        worktree_root: PathBuf,
1035        subroot_relative_path: Arc<RelPath>,
1036        project_env: Option<HashMap<String, String>>,
1037        fs: &dyn Fs,
1038    ) -> ToolchainList {
1039        let env = project_env.unwrap_or_default();
1040        let environment = EnvironmentApi::from_env(&env);
1041        let locators = pet::locators::create_locators(
1042            Arc::new(pet_conda::Conda::from(&environment)),
1043            Arc::new(pet_poetry::Poetry::from(&environment)),
1044            &environment,
1045        );
1046        let mut config = Configuration::default();
1047
1048        // `.ancestors()` will yield at least one path, so in case of empty `subroot_relative_path`, we'll just use
1049        // worktree root as the workspace directory.
1050        config.workspace_directories = Some(
1051            subroot_relative_path
1052                .ancestors()
1053                .map(|ancestor| worktree_root.join(ancestor.as_std_path()))
1054                .collect(),
1055        );
1056        for locator in locators.iter() {
1057            locator.configure(&config);
1058        }
1059
1060        let reporter = pet_reporter::collect::create_reporter();
1061        pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
1062
1063        let mut toolchains = reporter
1064            .environments
1065            .lock()
1066            .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
1067
1068        let wr = worktree_root;
1069        let wr_venv = get_worktree_venv_declaration(&wr).await;
1070        // Sort detected environments by:
1071        //     environment name matching activation file (<workdir>/.venv)
1072        //     environment project dir matching worktree_root
1073        //     general env priority
1074        //     environment path matching the CONDA_PREFIX env var
1075        //     executable path
1076        toolchains.sort_by(|lhs, rhs| {
1077            // Compare venv names against worktree .venv file
1078            let venv_ordering =
1079                wr_venv
1080                    .as_ref()
1081                    .map_or(Ordering::Equal, |venv| match (&lhs.name, &rhs.name) {
1082                        (Some(l), Some(r)) => (r == venv).cmp(&(l == venv)),
1083                        (Some(l), None) if l == venv => Ordering::Less,
1084                        (None, Some(r)) if r == venv => Ordering::Greater,
1085                        _ => Ordering::Equal,
1086                    });
1087
1088            // Compare project paths against worktree root
1089            let proj_ordering = || {
1090                let lhs_project = lhs.project.clone().or_else(|| get_venv_parent_dir(lhs));
1091                let rhs_project = rhs.project.clone().or_else(|| get_venv_parent_dir(rhs));
1092                wr_distance(&wr, lhs_project.as_ref()).cmp(&wr_distance(&wr, rhs_project.as_ref()))
1093            };
1094
1095            // Compare environment priorities
1096            let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind));
1097
1098            // Compare conda prefixes
1099            let conda_ordering = || {
1100                if lhs.kind == Some(PythonEnvironmentKind::Conda) {
1101                    environment
1102                        .get_env_var("CONDA_PREFIX".to_string())
1103                        .map(|conda_prefix| {
1104                            let is_match = |exe: &Option<PathBuf>| {
1105                                exe.as_ref().is_some_and(|e| e.starts_with(&conda_prefix))
1106                            };
1107                            match (is_match(&lhs.executable), is_match(&rhs.executable)) {
1108                                (true, false) => Ordering::Less,
1109                                (false, true) => Ordering::Greater,
1110                                _ => Ordering::Equal,
1111                            }
1112                        })
1113                        .unwrap_or(Ordering::Equal)
1114                } else {
1115                    Ordering::Equal
1116                }
1117            };
1118
1119            // Compare Python executables
1120            let exe_ordering = || lhs.executable.cmp(&rhs.executable);
1121
1122            venv_ordering
1123                .then_with(proj_ordering)
1124                .then_with(priority_ordering)
1125                .then_with(conda_ordering)
1126                .then_with(exe_ordering)
1127        });
1128
1129        let mut out_toolchains = Vec::new();
1130        for toolchain in toolchains {
1131            let Some(toolchain) = venv_to_toolchain(toolchain, fs).await else {
1132                continue;
1133            };
1134            out_toolchains.push(toolchain);
1135        }
1136        out_toolchains.dedup();
1137        ToolchainList {
1138            toolchains: out_toolchains,
1139            default: None,
1140            groups: Default::default(),
1141        }
1142    }
1143    fn meta(&self) -> ToolchainMetadata {
1144        ToolchainMetadata {
1145            term: SharedString::new_static("Virtual Environment"),
1146            new_toolchain_placeholder: SharedString::new_static(
1147                "A path to the python3 executable within a virtual environment, or path to virtual environment itself",
1148            ),
1149            manifest_name: ManifestName::from(SharedString::new_static("pyproject.toml")),
1150        }
1151    }
1152
1153    async fn resolve(
1154        &self,
1155        path: PathBuf,
1156        env: Option<HashMap<String, String>>,
1157        fs: &dyn Fs,
1158    ) -> anyhow::Result<Toolchain> {
1159        let env = env.unwrap_or_default();
1160        let environment = EnvironmentApi::from_env(&env);
1161        let locators = pet::locators::create_locators(
1162            Arc::new(pet_conda::Conda::from(&environment)),
1163            Arc::new(pet_poetry::Poetry::from(&environment)),
1164            &environment,
1165        );
1166        let toolchain = pet::resolve::resolve_environment(&path, &locators, &environment)
1167            .context("Could not find a virtual environment in provided path")?;
1168        let venv = toolchain.resolved.unwrap_or(toolchain.discovered);
1169        venv_to_toolchain(venv, fs)
1170            .await
1171            .context("Could not convert a venv into a toolchain")
1172    }
1173
1174    fn activation_script(&self, toolchain: &Toolchain, shell: ShellKind) -> Vec<String> {
1175        let Ok(toolchain) =
1176            serde_json::from_value::<PythonToolchainData>(toolchain.as_json.clone())
1177        else {
1178            return vec![];
1179        };
1180
1181        log::debug!("(Python) Composing activation script for toolchain {toolchain:?}");
1182
1183        let mut activation_script = vec![];
1184
1185        match toolchain.environment.kind {
1186            Some(PythonEnvironmentKind::Conda) => {
1187                if let Some(name) = &toolchain.environment.name {
1188                    activation_script.push(format!("conda activate {name}"));
1189                } else {
1190                    activation_script.push("conda activate".to_string());
1191                }
1192            }
1193            Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
1194                if let Some(activation_scripts) = &toolchain.activation_scripts {
1195                    if let Some(activate_script_path) = activation_scripts.get(&shell) {
1196                        let activate_keyword = shell.activate_keyword();
1197                        if let Some(quoted) =
1198                            shell.try_quote(&activate_script_path.to_string_lossy())
1199                        {
1200                            activation_script.push(format!("{activate_keyword} {quoted}"));
1201                        }
1202                    }
1203                }
1204            }
1205            Some(PythonEnvironmentKind::Pyenv) => {
1206                let Some(manager) = &toolchain.environment.manager else {
1207                    return vec![];
1208                };
1209                let version = toolchain.environment.version.as_deref().unwrap_or("system");
1210                let pyenv = &manager.executable;
1211                let pyenv = pyenv.display();
1212                activation_script.extend(match shell {
1213                    ShellKind::Fish => Some(format!("\"{pyenv}\" shell - fish {version}")),
1214                    ShellKind::Posix => Some(format!("\"{pyenv}\" shell - sh {version}")),
1215                    ShellKind::Nushell => Some(format!("^\"{pyenv}\" shell - nu {version}")),
1216                    ShellKind::PowerShell => None,
1217                    ShellKind::Csh => None,
1218                    ShellKind::Tcsh => None,
1219                    ShellKind::Cmd => None,
1220                    ShellKind::Rc => None,
1221                    ShellKind::Xonsh => None,
1222                    ShellKind::Elvish => None,
1223                })
1224            }
1225            _ => {}
1226        }
1227        activation_script
1228    }
1229}
1230
1231async fn venv_to_toolchain(venv: PythonEnvironment, fs: &dyn Fs) -> Option<Toolchain> {
1232    let mut name = String::from("Python");
1233    if let Some(ref version) = venv.version {
1234        _ = write!(name, " {version}");
1235    }
1236
1237    let name_and_kind = match (&venv.name, &venv.kind) {
1238        (Some(name), Some(kind)) => Some(format!("({name}; {})", python_env_kind_display(kind))),
1239        (Some(name), None) => Some(format!("({name})")),
1240        (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
1241        (None, None) => None,
1242    };
1243
1244    if let Some(nk) = name_and_kind {
1245        _ = write!(name, " {nk}");
1246    }
1247
1248    let mut activation_scripts = HashMap::default();
1249    match venv.kind {
1250        Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
1251            resolve_venv_activation_scripts(&venv, fs, &mut activation_scripts).await
1252        }
1253        _ => {}
1254    }
1255    let data = PythonToolchainData {
1256        environment: venv,
1257        activation_scripts: Some(activation_scripts),
1258    };
1259
1260    Some(Toolchain {
1261        name: name.into(),
1262        path: data
1263            .environment
1264            .executable
1265            .as_ref()?
1266            .to_str()?
1267            .to_owned()
1268            .into(),
1269        language_name: LanguageName::new("Python"),
1270        as_json: serde_json::to_value(data).ok()?,
1271    })
1272}
1273
1274async fn resolve_venv_activation_scripts(
1275    venv: &PythonEnvironment,
1276    fs: &dyn Fs,
1277    activation_scripts: &mut HashMap<ShellKind, PathBuf>,
1278) {
1279    log::debug!("(Python) Resolving activation scripts for venv toolchain {venv:?}");
1280    if let Some(prefix) = &venv.prefix {
1281        for (shell_kind, script_name) in &[
1282            (ShellKind::Posix, "activate"),
1283            (ShellKind::Rc, "activate"),
1284            (ShellKind::Csh, "activate.csh"),
1285            (ShellKind::Tcsh, "activate.csh"),
1286            (ShellKind::Fish, "activate.fish"),
1287            (ShellKind::Nushell, "activate.nu"),
1288            (ShellKind::PowerShell, "activate.ps1"),
1289            (ShellKind::Cmd, "activate.bat"),
1290            (ShellKind::Xonsh, "activate.xsh"),
1291        ] {
1292            let path = prefix.join(BINARY_DIR).join(script_name);
1293
1294            log::debug!("Trying path: {}", path.display());
1295
1296            if fs.is_file(&path).await {
1297                activation_scripts.insert(*shell_kind, path);
1298            }
1299        }
1300    }
1301}
1302
1303pub struct EnvironmentApi<'a> {
1304    global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
1305    project_env: &'a HashMap<String, String>,
1306    pet_env: pet_core::os_environment::EnvironmentApi,
1307}
1308
1309impl<'a> EnvironmentApi<'a> {
1310    pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
1311        let paths = project_env
1312            .get("PATH")
1313            .map(|p| std::env::split_paths(p).collect())
1314            .unwrap_or_default();
1315
1316        EnvironmentApi {
1317            global_search_locations: Arc::new(Mutex::new(paths)),
1318            project_env,
1319            pet_env: pet_core::os_environment::EnvironmentApi::new(),
1320        }
1321    }
1322
1323    fn user_home(&self) -> Option<PathBuf> {
1324        self.project_env
1325            .get("HOME")
1326            .or_else(|| self.project_env.get("USERPROFILE"))
1327            .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
1328            .or_else(|| self.pet_env.get_user_home())
1329    }
1330}
1331
1332impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
1333    fn get_user_home(&self) -> Option<PathBuf> {
1334        self.user_home()
1335    }
1336
1337    fn get_root(&self) -> Option<PathBuf> {
1338        None
1339    }
1340
1341    fn get_env_var(&self, key: String) -> Option<String> {
1342        self.project_env
1343            .get(&key)
1344            .cloned()
1345            .or_else(|| self.pet_env.get_env_var(key))
1346    }
1347
1348    fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
1349        if self.global_search_locations.lock().is_empty() {
1350            let mut paths = std::env::split_paths(
1351                &self
1352                    .get_env_var("PATH".to_string())
1353                    .or_else(|| self.get_env_var("Path".to_string()))
1354                    .unwrap_or_default(),
1355            )
1356            .collect::<Vec<PathBuf>>();
1357
1358            log::trace!("Env PATH: {:?}", paths);
1359            for p in self.pet_env.get_know_global_search_locations() {
1360                if !paths.contains(&p) {
1361                    paths.push(p);
1362                }
1363            }
1364
1365            let mut paths = paths
1366                .into_iter()
1367                .filter(|p| p.exists())
1368                .collect::<Vec<PathBuf>>();
1369
1370            self.global_search_locations.lock().append(&mut paths);
1371        }
1372        self.global_search_locations.lock().clone()
1373    }
1374}
1375
1376pub(crate) struct PyLspAdapter {
1377    python_venv_base: OnceCell<Result<Arc<Path>, String>>,
1378}
1379impl PyLspAdapter {
1380    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
1381    pub(crate) fn new() -> Self {
1382        Self {
1383            python_venv_base: OnceCell::new(),
1384        }
1385    }
1386    async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
1387        let python_path = Self::find_base_python(delegate)
1388            .await
1389            .with_context(|| {
1390                let mut message = "Could not find Python installation for PyLSP".to_owned();
1391                if cfg!(windows){
1392                    message.push_str(". Install Python from the Microsoft Store, or manually from https://www.python.org/downloads/windows.")
1393                }
1394                message
1395            })?;
1396        let work_dir = delegate
1397            .language_server_download_dir(&Self::SERVER_NAME)
1398            .await
1399            .context("Could not get working directory for PyLSP")?;
1400        let mut path = PathBuf::from(work_dir.as_ref());
1401        path.push("pylsp-venv");
1402        if !path.exists() {
1403            util::command::new_smol_command(python_path)
1404                .arg("-m")
1405                .arg("venv")
1406                .arg("pylsp-venv")
1407                .current_dir(work_dir)
1408                .spawn()?
1409                .output()
1410                .await?;
1411        }
1412
1413        Ok(path.into())
1414    }
1415    // Find "baseline", user python version from which we'll create our own venv.
1416    async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
1417        for path in ["python3", "python"] {
1418            let Some(path) = delegate.which(path.as_ref()).await else {
1419                continue;
1420            };
1421            // Try to detect situations where `python3` exists but is not a real Python interpreter.
1422            // Notably, on fresh Windows installs, `python3` is a shim that opens the Microsoft Store app
1423            // when run with no arguments, and just fails otherwise.
1424            let Some(output) = new_smol_command(&path)
1425                .args(["-c", "print(1 + 2)"])
1426                .output()
1427                .await
1428                .ok()
1429            else {
1430                continue;
1431            };
1432            if output.stdout.trim_ascii() != b"3" {
1433                continue;
1434            }
1435            return Some(path);
1436        }
1437        None
1438    }
1439
1440    async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
1441        self.python_venv_base
1442            .get_or_init(move || async move {
1443                Self::ensure_venv(delegate)
1444                    .await
1445                    .map_err(|e| format!("{e}"))
1446            })
1447            .await
1448            .clone()
1449    }
1450}
1451
1452const BINARY_DIR: &str = if cfg!(target_os = "windows") {
1453    "Scripts"
1454} else {
1455    "bin"
1456};
1457
1458#[async_trait(?Send)]
1459impl LspAdapter for PyLspAdapter {
1460    fn name(&self) -> LanguageServerName {
1461        Self::SERVER_NAME
1462    }
1463
1464    async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
1465
1466    async fn label_for_completion(
1467        &self,
1468        item: &lsp::CompletionItem,
1469        language: &Arc<language::Language>,
1470    ) -> Option<language::CodeLabel> {
1471        let label = &item.label;
1472        let label_len = label.len();
1473        let grammar = language.grammar()?;
1474        let highlight_id = match item.kind? {
1475            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1476            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1477            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1478            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1479            _ => return None,
1480        };
1481        Some(language::CodeLabel::filtered(
1482            label.clone(),
1483            label_len,
1484            item.filter_text.as_deref(),
1485            vec![(0..label.len(), highlight_id)],
1486        ))
1487    }
1488
1489    async fn label_for_symbol(
1490        &self,
1491        name: &str,
1492        kind: lsp::SymbolKind,
1493        language: &Arc<language::Language>,
1494    ) -> Option<language::CodeLabel> {
1495        let (text, filter_range, display_range) = match kind {
1496            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1497                let text = format!("def {}():\n", name);
1498                let filter_range = 4..4 + name.len();
1499                let display_range = 0..filter_range.end;
1500                (text, filter_range, display_range)
1501            }
1502            lsp::SymbolKind::CLASS => {
1503                let text = format!("class {}:", name);
1504                let filter_range = 6..6 + name.len();
1505                let display_range = 0..filter_range.end;
1506                (text, filter_range, display_range)
1507            }
1508            lsp::SymbolKind::CONSTANT => {
1509                let text = format!("{} = 0", name);
1510                let filter_range = 0..name.len();
1511                let display_range = 0..filter_range.end;
1512                (text, filter_range, display_range)
1513            }
1514            _ => return None,
1515        };
1516        Some(language::CodeLabel::new(
1517            text[display_range.clone()].to_string(),
1518            filter_range,
1519            language.highlight_text(&text.as_str().into(), display_range),
1520        ))
1521    }
1522
1523    async fn workspace_configuration(
1524        self: Arc<Self>,
1525        adapter: &Arc<dyn LspAdapterDelegate>,
1526        toolchain: Option<Toolchain>,
1527        cx: &mut AsyncApp,
1528    ) -> Result<Value> {
1529        cx.update(move |cx| {
1530            let mut user_settings =
1531                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1532                    .and_then(|s| s.settings.clone())
1533                    .unwrap_or_else(|| {
1534                        json!({
1535                            "plugins": {
1536                                "pycodestyle": {"enabled": false},
1537                                "rope_autoimport": {"enabled": true, "memory": true},
1538                                "pylsp_mypy": {"enabled": false}
1539                            },
1540                            "rope": {
1541                                "ropeFolder": null
1542                            },
1543                        })
1544                    });
1545
1546            // If user did not explicitly modify their python venv, use one from picker.
1547            if let Some(toolchain) = toolchain {
1548                if !user_settings.is_object() {
1549                    user_settings = Value::Object(serde_json::Map::default());
1550                }
1551                let object = user_settings.as_object_mut().unwrap();
1552                if let Some(python) = object
1553                    .entry("plugins")
1554                    .or_insert(Value::Object(serde_json::Map::default()))
1555                    .as_object_mut()
1556                {
1557                    if let Some(jedi) = python
1558                        .entry("jedi")
1559                        .or_insert(Value::Object(serde_json::Map::default()))
1560                        .as_object_mut()
1561                    {
1562                        jedi.entry("environment".to_string())
1563                            .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1564                    }
1565                    if let Some(pylint) = python
1566                        .entry("pylsp_mypy")
1567                        .or_insert(Value::Object(serde_json::Map::default()))
1568                        .as_object_mut()
1569                    {
1570                        pylint.entry("overrides".to_string()).or_insert_with(|| {
1571                            Value::Array(vec![
1572                                Value::String("--python-executable".into()),
1573                                Value::String(toolchain.path.into()),
1574                                Value::String("--cache-dir=/dev/null".into()),
1575                                Value::Bool(true),
1576                            ])
1577                        });
1578                    }
1579                }
1580            }
1581            user_settings = Value::Object(serde_json::Map::from_iter([(
1582                "pylsp".to_string(),
1583                user_settings,
1584            )]));
1585
1586            user_settings
1587        })
1588    }
1589}
1590
1591impl LspInstaller for PyLspAdapter {
1592    type BinaryVersion = ();
1593    async fn check_if_user_installed(
1594        &self,
1595        delegate: &dyn LspAdapterDelegate,
1596        toolchain: Option<Toolchain>,
1597        _: &AsyncApp,
1598    ) -> Option<LanguageServerBinary> {
1599        if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
1600            let env = delegate.shell_env().await;
1601            Some(LanguageServerBinary {
1602                path: pylsp_bin,
1603                env: Some(env),
1604                arguments: vec![],
1605            })
1606        } else {
1607            let toolchain = toolchain?;
1608            let pylsp_path = Path::new(toolchain.path.as_ref()).parent()?.join("pylsp");
1609            pylsp_path.exists().then(|| LanguageServerBinary {
1610                path: toolchain.path.to_string().into(),
1611                arguments: vec![pylsp_path.into()],
1612                env: None,
1613            })
1614        }
1615    }
1616
1617    async fn fetch_latest_server_version(
1618        &self,
1619        _: &dyn LspAdapterDelegate,
1620        _: bool,
1621        _: &mut AsyncApp,
1622    ) -> Result<()> {
1623        Ok(())
1624    }
1625
1626    async fn fetch_server_binary(
1627        &self,
1628        _: (),
1629        _: PathBuf,
1630        delegate: &dyn LspAdapterDelegate,
1631    ) -> Result<LanguageServerBinary> {
1632        let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
1633        let pip_path = venv.join(BINARY_DIR).join("pip3");
1634        ensure!(
1635            util::command::new_smol_command(pip_path.as_path())
1636                .arg("install")
1637                .arg("python-lsp-server[all]")
1638                .arg("--upgrade")
1639                .output()
1640                .await?
1641                .status
1642                .success(),
1643            "python-lsp-server[all] installation failed"
1644        );
1645        ensure!(
1646            util::command::new_smol_command(pip_path)
1647                .arg("install")
1648                .arg("pylsp-mypy")
1649                .arg("--upgrade")
1650                .output()
1651                .await?
1652                .status
1653                .success(),
1654            "pylsp-mypy installation failed"
1655        );
1656        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1657        ensure!(
1658            delegate.which(pylsp.as_os_str()).await.is_some(),
1659            "pylsp installation was incomplete"
1660        );
1661        Ok(LanguageServerBinary {
1662            path: pylsp,
1663            env: None,
1664            arguments: vec![],
1665        })
1666    }
1667
1668    async fn cached_server_binary(
1669        &self,
1670        _: PathBuf,
1671        delegate: &dyn LspAdapterDelegate,
1672    ) -> Option<LanguageServerBinary> {
1673        let venv = self.base_venv(delegate).await.ok()?;
1674        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1675        delegate.which(pylsp.as_os_str()).await?;
1676        Some(LanguageServerBinary {
1677            path: pylsp,
1678            env: None,
1679            arguments: vec![],
1680        })
1681    }
1682}
1683
1684pub(crate) struct BasedPyrightLspAdapter {
1685    node: NodeRuntime,
1686}
1687
1688impl BasedPyrightLspAdapter {
1689    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("basedpyright");
1690    const BINARY_NAME: &'static str = "basedpyright-langserver";
1691    const SERVER_PATH: &str = "node_modules/basedpyright/langserver.index.js";
1692    const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "basedpyright/langserver.index.js";
1693
1694    pub(crate) fn new(node: NodeRuntime) -> Self {
1695        BasedPyrightLspAdapter { node }
1696    }
1697
1698    async fn get_cached_server_binary(
1699        container_dir: PathBuf,
1700        node: &NodeRuntime,
1701    ) -> Option<LanguageServerBinary> {
1702        let server_path = container_dir.join(Self::SERVER_PATH);
1703        if server_path.exists() {
1704            Some(LanguageServerBinary {
1705                path: node.binary_path().await.log_err()?,
1706                env: None,
1707                arguments: vec![server_path.into(), "--stdio".into()],
1708            })
1709        } else {
1710            log::error!("missing executable in directory {:?}", server_path);
1711            None
1712        }
1713    }
1714}
1715
1716#[async_trait(?Send)]
1717impl LspAdapter for BasedPyrightLspAdapter {
1718    fn name(&self) -> LanguageServerName {
1719        Self::SERVER_NAME
1720    }
1721
1722    async fn initialization_options(
1723        self: Arc<Self>,
1724        _: &Arc<dyn LspAdapterDelegate>,
1725    ) -> Result<Option<Value>> {
1726        // Provide minimal initialization options
1727        // Virtual environment configuration will be handled through workspace configuration
1728        Ok(Some(json!({
1729            "python": {
1730                "analysis": {
1731                    "autoSearchPaths": true,
1732                    "useLibraryCodeForTypes": true,
1733                    "autoImportCompletions": true
1734                }
1735            }
1736        })))
1737    }
1738
1739    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
1740        process_pyright_completions(items);
1741    }
1742
1743    async fn label_for_completion(
1744        &self,
1745        item: &lsp::CompletionItem,
1746        language: &Arc<language::Language>,
1747    ) -> Option<language::CodeLabel> {
1748        let label = &item.label;
1749        let label_len = label.len();
1750        let grammar = language.grammar()?;
1751        let highlight_id = match item.kind? {
1752            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method"),
1753            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function"),
1754            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type"),
1755            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant"),
1756            lsp::CompletionItemKind::VARIABLE => grammar.highlight_id_for_name("variable"),
1757            _ => {
1758                return None;
1759            }
1760        };
1761        let mut text = label.clone();
1762        if let Some(completion_details) = item
1763            .label_details
1764            .as_ref()
1765            .and_then(|details| details.description.as_ref())
1766        {
1767            write!(&mut text, " {}", completion_details).ok();
1768        }
1769        Some(language::CodeLabel::filtered(
1770            text,
1771            label_len,
1772            item.filter_text.as_deref(),
1773            highlight_id
1774                .map(|id| (0..label.len(), id))
1775                .into_iter()
1776                .collect(),
1777        ))
1778    }
1779
1780    async fn label_for_symbol(
1781        &self,
1782        name: &str,
1783        kind: lsp::SymbolKind,
1784        language: &Arc<language::Language>,
1785    ) -> Option<language::CodeLabel> {
1786        let (text, filter_range, display_range) = match kind {
1787            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1788                let text = format!("def {}():\n", name);
1789                let filter_range = 4..4 + name.len();
1790                let display_range = 0..filter_range.end;
1791                (text, filter_range, display_range)
1792            }
1793            lsp::SymbolKind::CLASS => {
1794                let text = format!("class {}:", name);
1795                let filter_range = 6..6 + name.len();
1796                let display_range = 0..filter_range.end;
1797                (text, filter_range, display_range)
1798            }
1799            lsp::SymbolKind::CONSTANT => {
1800                let text = format!("{} = 0", name);
1801                let filter_range = 0..name.len();
1802                let display_range = 0..filter_range.end;
1803                (text, filter_range, display_range)
1804            }
1805            _ => return None,
1806        };
1807        Some(language::CodeLabel::new(
1808            text[display_range.clone()].to_string(),
1809            filter_range,
1810            language.highlight_text(&text.as_str().into(), display_range),
1811        ))
1812    }
1813
1814    async fn workspace_configuration(
1815        self: Arc<Self>,
1816        adapter: &Arc<dyn LspAdapterDelegate>,
1817        toolchain: Option<Toolchain>,
1818        cx: &mut AsyncApp,
1819    ) -> Result<Value> {
1820        cx.update(move |cx| {
1821            let mut user_settings =
1822                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1823                    .and_then(|s| s.settings.clone())
1824                    .unwrap_or_default();
1825
1826            // If we have a detected toolchain, configure Pyright to use it
1827            if let Some(toolchain) = toolchain
1828                && let Ok(env) = serde_json::from_value::<
1829                    pet_core::python_environment::PythonEnvironment,
1830                >(toolchain.as_json.clone())
1831            {
1832                if !user_settings.is_object() {
1833                    user_settings = Value::Object(serde_json::Map::default());
1834                }
1835                let object = user_settings.as_object_mut().unwrap();
1836
1837                let interpreter_path = toolchain.path.to_string();
1838                if let Some(venv_dir) = env.prefix {
1839                    // Set venvPath and venv at the root level
1840                    // This matches the format of a pyrightconfig.json file
1841                    if let Some(parent) = venv_dir.parent() {
1842                        // Use relative path if the venv is inside the workspace
1843                        let venv_path = if parent == adapter.worktree_root_path() {
1844                            ".".to_string()
1845                        } else {
1846                            parent.to_string_lossy().into_owned()
1847                        };
1848                        object.insert("venvPath".to_string(), Value::String(venv_path));
1849                    }
1850
1851                    if let Some(venv_name) = venv_dir.file_name() {
1852                        object.insert(
1853                            "venv".to_owned(),
1854                            Value::String(venv_name.to_string_lossy().into_owned()),
1855                        );
1856                    }
1857                }
1858
1859                // Set both pythonPath and defaultInterpreterPath for compatibility
1860                if let Some(python) = object
1861                    .entry("python")
1862                    .or_insert(Value::Object(serde_json::Map::default()))
1863                    .as_object_mut()
1864                {
1865                    python.insert(
1866                        "pythonPath".to_owned(),
1867                        Value::String(interpreter_path.clone()),
1868                    );
1869                    python.insert(
1870                        "defaultInterpreterPath".to_owned(),
1871                        Value::String(interpreter_path),
1872                    );
1873                }
1874                // Basedpyright by default uses `strict` type checking, we tone it down as to not surpris users
1875                maybe!({
1876                    let analysis = object
1877                        .entry("basedpyright.analysis")
1878                        .or_insert(Value::Object(serde_json::Map::default()));
1879                    if let serde_json::map::Entry::Vacant(v) =
1880                        analysis.as_object_mut()?.entry("typeCheckingMode")
1881                    {
1882                        v.insert(Value::String("standard".to_owned()));
1883                    }
1884                    Some(())
1885                });
1886            }
1887
1888            user_settings
1889        })
1890    }
1891}
1892
1893impl LspInstaller for BasedPyrightLspAdapter {
1894    type BinaryVersion = String;
1895
1896    async fn fetch_latest_server_version(
1897        &self,
1898        _: &dyn LspAdapterDelegate,
1899        _: bool,
1900        _: &mut AsyncApp,
1901    ) -> Result<String> {
1902        self.node
1903            .npm_package_latest_version(Self::SERVER_NAME.as_ref())
1904            .await
1905    }
1906
1907    async fn check_if_user_installed(
1908        &self,
1909        delegate: &dyn LspAdapterDelegate,
1910        _: Option<Toolchain>,
1911        _: &AsyncApp,
1912    ) -> Option<LanguageServerBinary> {
1913        if let Some(path) = delegate.which(Self::BINARY_NAME.as_ref()).await {
1914            let env = delegate.shell_env().await;
1915            Some(LanguageServerBinary {
1916                path,
1917                env: Some(env),
1918                arguments: vec!["--stdio".into()],
1919            })
1920        } else {
1921            // TODO shouldn't this be self.node.binary_path()?
1922            let node = delegate.which("node".as_ref()).await?;
1923            let (node_modules_path, _) = delegate
1924                .npm_package_installed_version(Self::SERVER_NAME.as_ref())
1925                .await
1926                .log_err()??;
1927
1928            let path = node_modules_path.join(Self::NODE_MODULE_RELATIVE_SERVER_PATH);
1929
1930            let env = delegate.shell_env().await;
1931            Some(LanguageServerBinary {
1932                path: node,
1933                env: Some(env),
1934                arguments: vec![path.into(), "--stdio".into()],
1935            })
1936        }
1937    }
1938
1939    async fn fetch_server_binary(
1940        &self,
1941        latest_version: Self::BinaryVersion,
1942        container_dir: PathBuf,
1943        delegate: &dyn LspAdapterDelegate,
1944    ) -> Result<LanguageServerBinary> {
1945        let server_path = container_dir.join(Self::SERVER_PATH);
1946
1947        self.node
1948            .npm_install_packages(
1949                &container_dir,
1950                &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
1951            )
1952            .await?;
1953
1954        let env = delegate.shell_env().await;
1955        Ok(LanguageServerBinary {
1956            path: self.node.binary_path().await?,
1957            env: Some(env),
1958            arguments: vec![server_path.into(), "--stdio".into()],
1959        })
1960    }
1961
1962    async fn check_if_version_installed(
1963        &self,
1964        version: &Self::BinaryVersion,
1965        container_dir: &PathBuf,
1966        delegate: &dyn LspAdapterDelegate,
1967    ) -> Option<LanguageServerBinary> {
1968        let server_path = container_dir.join(Self::SERVER_PATH);
1969
1970        let should_install_language_server = self
1971            .node
1972            .should_install_npm_package(
1973                Self::SERVER_NAME.as_ref(),
1974                &server_path,
1975                container_dir,
1976                VersionStrategy::Latest(version),
1977            )
1978            .await;
1979
1980        if should_install_language_server {
1981            None
1982        } else {
1983            let env = delegate.shell_env().await;
1984            Some(LanguageServerBinary {
1985                path: self.node.binary_path().await.ok()?,
1986                env: Some(env),
1987                arguments: vec![server_path.into(), "--stdio".into()],
1988            })
1989        }
1990    }
1991
1992    async fn cached_server_binary(
1993        &self,
1994        container_dir: PathBuf,
1995        delegate: &dyn LspAdapterDelegate,
1996    ) -> Option<LanguageServerBinary> {
1997        let mut binary = Self::get_cached_server_binary(container_dir, &self.node).await?;
1998        binary.env = Some(delegate.shell_env().await);
1999        Some(binary)
2000    }
2001}
2002
2003pub(crate) struct RuffLspAdapter {
2004    fs: Arc<dyn Fs>,
2005}
2006
2007#[cfg(target_os = "macos")]
2008impl RuffLspAdapter {
2009    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2010    const ARCH_SERVER_NAME: &str = "apple-darwin";
2011}
2012
2013#[cfg(target_os = "linux")]
2014impl RuffLspAdapter {
2015    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2016    const ARCH_SERVER_NAME: &str = "unknown-linux-gnu";
2017}
2018
2019#[cfg(target_os = "freebsd")]
2020impl RuffLspAdapter {
2021    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2022    const ARCH_SERVER_NAME: &str = "unknown-freebsd";
2023}
2024
2025#[cfg(target_os = "windows")]
2026impl RuffLspAdapter {
2027    const GITHUB_ASSET_KIND: AssetKind = AssetKind::Zip;
2028    const ARCH_SERVER_NAME: &str = "pc-windows-msvc";
2029}
2030
2031impl RuffLspAdapter {
2032    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("ruff");
2033
2034    pub fn new(fs: Arc<dyn Fs>) -> RuffLspAdapter {
2035        RuffLspAdapter { fs }
2036    }
2037
2038    fn build_asset_name() -> Result<(String, String)> {
2039        let arch = match consts::ARCH {
2040            "x86" => "i686",
2041            _ => consts::ARCH,
2042        };
2043        let os = Self::ARCH_SERVER_NAME;
2044        let suffix = match consts::OS {
2045            "windows" => "zip",
2046            _ => "tar.gz",
2047        };
2048        let asset_name = format!("ruff-{arch}-{os}.{suffix}");
2049        let asset_stem = format!("ruff-{arch}-{os}");
2050        Ok((asset_stem, asset_name))
2051    }
2052}
2053
2054#[async_trait(?Send)]
2055impl LspAdapter for RuffLspAdapter {
2056    fn name(&self) -> LanguageServerName {
2057        Self::SERVER_NAME
2058    }
2059}
2060
2061impl LspInstaller for RuffLspAdapter {
2062    type BinaryVersion = GitHubLspBinaryVersion;
2063    async fn check_if_user_installed(
2064        &self,
2065        delegate: &dyn LspAdapterDelegate,
2066        toolchain: Option<Toolchain>,
2067        _: &AsyncApp,
2068    ) -> Option<LanguageServerBinary> {
2069        let ruff_in_venv = if let Some(toolchain) = toolchain
2070            && toolchain.language_name.as_ref() == "Python"
2071        {
2072            Path::new(toolchain.path.as_str())
2073                .parent()
2074                .map(|path| path.join("ruff"))
2075        } else {
2076            None
2077        };
2078
2079        for path in ruff_in_venv.into_iter().chain(["ruff".into()]) {
2080            if let Some(ruff_bin) = delegate.which(path.as_os_str()).await {
2081                let env = delegate.shell_env().await;
2082                return Some(LanguageServerBinary {
2083                    path: ruff_bin,
2084                    env: Some(env),
2085                    arguments: vec!["server".into()],
2086                });
2087            }
2088        }
2089
2090        None
2091    }
2092
2093    async fn fetch_latest_server_version(
2094        &self,
2095        delegate: &dyn LspAdapterDelegate,
2096        _: bool,
2097        _: &mut AsyncApp,
2098    ) -> Result<GitHubLspBinaryVersion> {
2099        let release =
2100            latest_github_release("astral-sh/ruff", true, false, delegate.http_client()).await?;
2101        let (_, asset_name) = Self::build_asset_name()?;
2102        let asset = release
2103            .assets
2104            .into_iter()
2105            .find(|asset| asset.name == asset_name)
2106            .with_context(|| format!("no asset found matching `{asset_name:?}`"))?;
2107        Ok(GitHubLspBinaryVersion {
2108            name: release.tag_name,
2109            url: asset.browser_download_url,
2110            digest: asset.digest,
2111        })
2112    }
2113
2114    async fn fetch_server_binary(
2115        &self,
2116        latest_version: GitHubLspBinaryVersion,
2117        container_dir: PathBuf,
2118        delegate: &dyn LspAdapterDelegate,
2119    ) -> Result<LanguageServerBinary> {
2120        let GitHubLspBinaryVersion {
2121            name,
2122            url,
2123            digest: expected_digest,
2124        } = latest_version;
2125        let destination_path = container_dir.join(format!("ruff-{name}"));
2126        let server_path = match Self::GITHUB_ASSET_KIND {
2127            AssetKind::TarGz | AssetKind::Gz => destination_path
2128                .join(Self::build_asset_name()?.0)
2129                .join("ruff"),
2130            AssetKind::Zip => destination_path.clone().join("ruff.exe"),
2131        };
2132
2133        let binary = LanguageServerBinary {
2134            path: server_path.clone(),
2135            env: None,
2136            arguments: vec!["server".into()],
2137        };
2138
2139        let metadata_path = destination_path.with_extension("metadata");
2140        let metadata = GithubBinaryMetadata::read_from_file(&metadata_path)
2141            .await
2142            .ok();
2143        if let Some(metadata) = metadata {
2144            let validity_check = async || {
2145                delegate
2146                    .try_exec(LanguageServerBinary {
2147                        path: server_path.clone(),
2148                        arguments: vec!["--version".into()],
2149                        env: None,
2150                    })
2151                    .await
2152                    .inspect_err(|err| {
2153                        log::warn!("Unable to run {server_path:?} asset, redownloading: {err}",)
2154                    })
2155            };
2156            if let (Some(actual_digest), Some(expected_digest)) =
2157                (&metadata.digest, &expected_digest)
2158            {
2159                if actual_digest == expected_digest {
2160                    if validity_check().await.is_ok() {
2161                        return Ok(binary);
2162                    }
2163                } else {
2164                    log::info!(
2165                        "SHA-256 mismatch for {destination_path:?} asset, downloading new asset. Expected: {expected_digest}, Got: {actual_digest}"
2166                    );
2167                }
2168            } else if validity_check().await.is_ok() {
2169                return Ok(binary);
2170            }
2171        }
2172
2173        download_server_binary(
2174            &*delegate.http_client(),
2175            &url,
2176            expected_digest.as_deref(),
2177            &destination_path,
2178            Self::GITHUB_ASSET_KIND,
2179        )
2180        .await?;
2181        make_file_executable(&server_path).await?;
2182        remove_matching(&container_dir, |path| path != destination_path).await;
2183        GithubBinaryMetadata::write_to_file(
2184            &GithubBinaryMetadata {
2185                metadata_version: 1,
2186                digest: expected_digest,
2187            },
2188            &metadata_path,
2189        )
2190        .await?;
2191
2192        Ok(LanguageServerBinary {
2193            path: server_path,
2194            env: None,
2195            arguments: vec!["server".into()],
2196        })
2197    }
2198
2199    async fn cached_server_binary(
2200        &self,
2201        container_dir: PathBuf,
2202        _: &dyn LspAdapterDelegate,
2203    ) -> Option<LanguageServerBinary> {
2204        maybe!(async {
2205            let mut last = None;
2206            let mut entries = self.fs.read_dir(&container_dir).await?;
2207            while let Some(entry) = entries.next().await {
2208                let path = entry?;
2209                if path.extension().is_some_and(|ext| ext == "metadata") {
2210                    continue;
2211                }
2212                last = Some(path);
2213            }
2214
2215            let path = last.context("no cached binary")?;
2216            let path = match Self::GITHUB_ASSET_KIND {
2217                AssetKind::TarGz | AssetKind::Gz => {
2218                    path.join(Self::build_asset_name()?.0).join("ruff")
2219                }
2220                AssetKind::Zip => path.join("ruff.exe"),
2221            };
2222
2223            anyhow::Ok(LanguageServerBinary {
2224                path,
2225                env: None,
2226                arguments: vec!["server".into()],
2227            })
2228        })
2229        .await
2230        .log_err()
2231    }
2232}
2233
2234#[cfg(test)]
2235mod tests {
2236    use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
2237    use language::{AutoindentMode, Buffer};
2238    use settings::SettingsStore;
2239    use std::num::NonZeroU32;
2240
2241    #[gpui::test]
2242    async fn test_python_autoindent(cx: &mut TestAppContext) {
2243        cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
2244        let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
2245        cx.update(|cx| {
2246            let test_settings = SettingsStore::test(cx);
2247            cx.set_global(test_settings);
2248            language::init(cx);
2249            cx.update_global::<SettingsStore, _>(|store, cx| {
2250                store.update_user_settings(cx, |s| {
2251                    s.project.all_languages.defaults.tab_size = NonZeroU32::new(2);
2252                });
2253            });
2254        });
2255
2256        cx.new(|cx| {
2257            let mut buffer = Buffer::local("", cx).with_language(language, cx);
2258            let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
2259                let ix = buffer.len();
2260                buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
2261            };
2262
2263            // indent after "def():"
2264            append(&mut buffer, "def a():\n", cx);
2265            assert_eq!(buffer.text(), "def a():\n  ");
2266
2267            // preserve indent after blank line
2268            append(&mut buffer, "\n  ", cx);
2269            assert_eq!(buffer.text(), "def a():\n  \n  ");
2270
2271            // indent after "if"
2272            append(&mut buffer, "if a:\n  ", cx);
2273            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    ");
2274
2275            // preserve indent after statement
2276            append(&mut buffer, "b()\n", cx);
2277            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    ");
2278
2279            // preserve indent after statement
2280            append(&mut buffer, "else", cx);
2281            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    else");
2282
2283            // dedent "else""
2284            append(&mut buffer, ":", cx);
2285            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n  else:");
2286
2287            // indent lines after else
2288            append(&mut buffer, "\n", cx);
2289            assert_eq!(
2290                buffer.text(),
2291                "def a():\n  \n  if a:\n    b()\n  else:\n    "
2292            );
2293
2294            // indent after an open paren. the closing paren is not indented
2295            // because there is another token before it on the same line.
2296            append(&mut buffer, "foo(\n1)", cx);
2297            assert_eq!(
2298                buffer.text(),
2299                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n      1)"
2300            );
2301
2302            // dedent the closing paren if it is shifted to the beginning of the line
2303            let argument_ix = buffer.text().find('1').unwrap();
2304            buffer.edit(
2305                [(argument_ix..argument_ix + 1, "")],
2306                Some(AutoindentMode::EachLine),
2307                cx,
2308            );
2309            assert_eq!(
2310                buffer.text(),
2311                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )"
2312            );
2313
2314            // preserve indent after the close paren
2315            append(&mut buffer, "\n", cx);
2316            assert_eq!(
2317                buffer.text(),
2318                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n    "
2319            );
2320
2321            // manually outdent the last line
2322            let end_whitespace_ix = buffer.len() - 4;
2323            buffer.edit(
2324                [(end_whitespace_ix..buffer.len(), "")],
2325                Some(AutoindentMode::EachLine),
2326                cx,
2327            );
2328            assert_eq!(
2329                buffer.text(),
2330                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n"
2331            );
2332
2333            // preserve the newly reduced indentation on the next newline
2334            append(&mut buffer, "\n", cx);
2335            assert_eq!(
2336                buffer.text(),
2337                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n\n"
2338            );
2339
2340            // reset to a for loop statement
2341            let statement = "for i in range(10):\n  print(i)\n";
2342            buffer.edit([(0..buffer.len(), statement)], None, cx);
2343
2344            // insert single line comment after each line
2345            let eol_ixs = statement
2346                .char_indices()
2347                .filter_map(|(ix, c)| if c == '\n' { Some(ix) } else { None })
2348                .collect::<Vec<usize>>();
2349            let editions = eol_ixs
2350                .iter()
2351                .enumerate()
2352                .map(|(i, &eol_ix)| (eol_ix..eol_ix, format!(" # comment {}", i + 1)))
2353                .collect::<Vec<(std::ops::Range<usize>, String)>>();
2354            buffer.edit(editions, Some(AutoindentMode::EachLine), cx);
2355            assert_eq!(
2356                buffer.text(),
2357                "for i in range(10): # comment 1\n  print(i) # comment 2\n"
2358            );
2359
2360            // reset to a simple if statement
2361            buffer.edit([(0..buffer.len(), "if a:\n  b(\n  )")], None, cx);
2362
2363            // dedent "else" on the line after a closing paren
2364            append(&mut buffer, "\n  else:\n", cx);
2365            assert_eq!(buffer.text(), "if a:\n  b(\n  )\nelse:\n  ");
2366
2367            buffer
2368        });
2369    }
2370}