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                })
1223            }
1224            _ => {}
1225        }
1226        activation_script
1227    }
1228}
1229
1230async fn venv_to_toolchain(venv: PythonEnvironment, fs: &dyn Fs) -> Option<Toolchain> {
1231    let mut name = String::from("Python");
1232    if let Some(ref version) = venv.version {
1233        _ = write!(name, " {version}");
1234    }
1235
1236    let name_and_kind = match (&venv.name, &venv.kind) {
1237        (Some(name), Some(kind)) => Some(format!("({name}; {})", python_env_kind_display(kind))),
1238        (Some(name), None) => Some(format!("({name})")),
1239        (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
1240        (None, None) => None,
1241    };
1242
1243    if let Some(nk) = name_and_kind {
1244        _ = write!(name, " {nk}");
1245    }
1246
1247    let mut activation_scripts = HashMap::default();
1248    match venv.kind {
1249        Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
1250            resolve_venv_activation_scripts(&venv, fs, &mut activation_scripts).await
1251        }
1252        _ => {}
1253    }
1254    let data = PythonToolchainData {
1255        environment: venv,
1256        activation_scripts: Some(activation_scripts),
1257    };
1258
1259    Some(Toolchain {
1260        name: name.into(),
1261        path: data
1262            .environment
1263            .executable
1264            .as_ref()?
1265            .to_str()?
1266            .to_owned()
1267            .into(),
1268        language_name: LanguageName::new("Python"),
1269        as_json: serde_json::to_value(data).ok()?,
1270    })
1271}
1272
1273async fn resolve_venv_activation_scripts(
1274    venv: &PythonEnvironment,
1275    fs: &dyn Fs,
1276    activation_scripts: &mut HashMap<ShellKind, PathBuf>,
1277) {
1278    log::debug!("(Python) Resolving activation scripts for venv toolchain {venv:?}");
1279    if let Some(prefix) = &venv.prefix {
1280        for (shell_kind, script_name) in &[
1281            (ShellKind::Posix, "activate"),
1282            (ShellKind::Rc, "activate"),
1283            (ShellKind::Csh, "activate.csh"),
1284            (ShellKind::Tcsh, "activate.csh"),
1285            (ShellKind::Fish, "activate.fish"),
1286            (ShellKind::Nushell, "activate.nu"),
1287            (ShellKind::PowerShell, "activate.ps1"),
1288            (ShellKind::Cmd, "activate.bat"),
1289            (ShellKind::Xonsh, "activate.xsh"),
1290        ] {
1291            let path = prefix.join(BINARY_DIR).join(script_name);
1292
1293            log::debug!("Trying path: {}", path.display());
1294
1295            if fs.is_file(&path).await {
1296                activation_scripts.insert(*shell_kind, path);
1297            }
1298        }
1299    }
1300}
1301
1302pub struct EnvironmentApi<'a> {
1303    global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
1304    project_env: &'a HashMap<String, String>,
1305    pet_env: pet_core::os_environment::EnvironmentApi,
1306}
1307
1308impl<'a> EnvironmentApi<'a> {
1309    pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
1310        let paths = project_env
1311            .get("PATH")
1312            .map(|p| std::env::split_paths(p).collect())
1313            .unwrap_or_default();
1314
1315        EnvironmentApi {
1316            global_search_locations: Arc::new(Mutex::new(paths)),
1317            project_env,
1318            pet_env: pet_core::os_environment::EnvironmentApi::new(),
1319        }
1320    }
1321
1322    fn user_home(&self) -> Option<PathBuf> {
1323        self.project_env
1324            .get("HOME")
1325            .or_else(|| self.project_env.get("USERPROFILE"))
1326            .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
1327            .or_else(|| self.pet_env.get_user_home())
1328    }
1329}
1330
1331impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
1332    fn get_user_home(&self) -> Option<PathBuf> {
1333        self.user_home()
1334    }
1335
1336    fn get_root(&self) -> Option<PathBuf> {
1337        None
1338    }
1339
1340    fn get_env_var(&self, key: String) -> Option<String> {
1341        self.project_env
1342            .get(&key)
1343            .cloned()
1344            .or_else(|| self.pet_env.get_env_var(key))
1345    }
1346
1347    fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
1348        if self.global_search_locations.lock().is_empty() {
1349            let mut paths = std::env::split_paths(
1350                &self
1351                    .get_env_var("PATH".to_string())
1352                    .or_else(|| self.get_env_var("Path".to_string()))
1353                    .unwrap_or_default(),
1354            )
1355            .collect::<Vec<PathBuf>>();
1356
1357            log::trace!("Env PATH: {:?}", paths);
1358            for p in self.pet_env.get_know_global_search_locations() {
1359                if !paths.contains(&p) {
1360                    paths.push(p);
1361                }
1362            }
1363
1364            let mut paths = paths
1365                .into_iter()
1366                .filter(|p| p.exists())
1367                .collect::<Vec<PathBuf>>();
1368
1369            self.global_search_locations.lock().append(&mut paths);
1370        }
1371        self.global_search_locations.lock().clone()
1372    }
1373}
1374
1375pub(crate) struct PyLspAdapter {
1376    python_venv_base: OnceCell<Result<Arc<Path>, String>>,
1377}
1378impl PyLspAdapter {
1379    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
1380    pub(crate) fn new() -> Self {
1381        Self {
1382            python_venv_base: OnceCell::new(),
1383        }
1384    }
1385    async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
1386        let python_path = Self::find_base_python(delegate)
1387            .await
1388            .with_context(|| {
1389                let mut message = "Could not find Python installation for PyLSP".to_owned();
1390                if cfg!(windows){
1391                    message.push_str(". Install Python from the Microsoft Store, or manually from https://www.python.org/downloads/windows.")
1392                }
1393                message
1394            })?;
1395        let work_dir = delegate
1396            .language_server_download_dir(&Self::SERVER_NAME)
1397            .await
1398            .context("Could not get working directory for PyLSP")?;
1399        let mut path = PathBuf::from(work_dir.as_ref());
1400        path.push("pylsp-venv");
1401        if !path.exists() {
1402            util::command::new_smol_command(python_path)
1403                .arg("-m")
1404                .arg("venv")
1405                .arg("pylsp-venv")
1406                .current_dir(work_dir)
1407                .spawn()?
1408                .output()
1409                .await?;
1410        }
1411
1412        Ok(path.into())
1413    }
1414    // Find "baseline", user python version from which we'll create our own venv.
1415    async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
1416        for path in ["python3", "python"] {
1417            let Some(path) = delegate.which(path.as_ref()).await else {
1418                continue;
1419            };
1420            // Try to detect situations where `python3` exists but is not a real Python interpreter.
1421            // Notably, on fresh Windows installs, `python3` is a shim that opens the Microsoft Store app
1422            // when run with no arguments, and just fails otherwise.
1423            let Some(output) = new_smol_command(&path)
1424                .args(["-c", "print(1 + 2)"])
1425                .output()
1426                .await
1427                .ok()
1428            else {
1429                continue;
1430            };
1431            if output.stdout.trim_ascii() != b"3" {
1432                continue;
1433            }
1434            return Some(path);
1435        }
1436        None
1437    }
1438
1439    async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
1440        self.python_venv_base
1441            .get_or_init(move || async move {
1442                Self::ensure_venv(delegate)
1443                    .await
1444                    .map_err(|e| format!("{e}"))
1445            })
1446            .await
1447            .clone()
1448    }
1449}
1450
1451const BINARY_DIR: &str = if cfg!(target_os = "windows") {
1452    "Scripts"
1453} else {
1454    "bin"
1455};
1456
1457#[async_trait(?Send)]
1458impl LspAdapter for PyLspAdapter {
1459    fn name(&self) -> LanguageServerName {
1460        Self::SERVER_NAME
1461    }
1462
1463    async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
1464
1465    async fn label_for_completion(
1466        &self,
1467        item: &lsp::CompletionItem,
1468        language: &Arc<language::Language>,
1469    ) -> Option<language::CodeLabel> {
1470        let label = &item.label;
1471        let label_len = label.len();
1472        let grammar = language.grammar()?;
1473        let highlight_id = match item.kind? {
1474            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1475            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1476            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1477            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1478            _ => return None,
1479        };
1480        Some(language::CodeLabel::filtered(
1481            label.clone(),
1482            label_len,
1483            item.filter_text.as_deref(),
1484            vec![(0..label.len(), highlight_id)],
1485        ))
1486    }
1487
1488    async fn label_for_symbol(
1489        &self,
1490        name: &str,
1491        kind: lsp::SymbolKind,
1492        language: &Arc<language::Language>,
1493    ) -> Option<language::CodeLabel> {
1494        let (text, filter_range, display_range) = match kind {
1495            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1496                let text = format!("def {}():\n", name);
1497                let filter_range = 4..4 + name.len();
1498                let display_range = 0..filter_range.end;
1499                (text, filter_range, display_range)
1500            }
1501            lsp::SymbolKind::CLASS => {
1502                let text = format!("class {}:", name);
1503                let filter_range = 6..6 + name.len();
1504                let display_range = 0..filter_range.end;
1505                (text, filter_range, display_range)
1506            }
1507            lsp::SymbolKind::CONSTANT => {
1508                let text = format!("{} = 0", name);
1509                let filter_range = 0..name.len();
1510                let display_range = 0..filter_range.end;
1511                (text, filter_range, display_range)
1512            }
1513            _ => return None,
1514        };
1515        Some(language::CodeLabel::new(
1516            text[display_range.clone()].to_string(),
1517            filter_range,
1518            language.highlight_text(&text.as_str().into(), display_range),
1519        ))
1520    }
1521
1522    async fn workspace_configuration(
1523        self: Arc<Self>,
1524        adapter: &Arc<dyn LspAdapterDelegate>,
1525        toolchain: Option<Toolchain>,
1526        cx: &mut AsyncApp,
1527    ) -> Result<Value> {
1528        cx.update(move |cx| {
1529            let mut user_settings =
1530                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1531                    .and_then(|s| s.settings.clone())
1532                    .unwrap_or_else(|| {
1533                        json!({
1534                            "plugins": {
1535                                "pycodestyle": {"enabled": false},
1536                                "rope_autoimport": {"enabled": true, "memory": true},
1537                                "pylsp_mypy": {"enabled": false}
1538                            },
1539                            "rope": {
1540                                "ropeFolder": null
1541                            },
1542                        })
1543                    });
1544
1545            // If user did not explicitly modify their python venv, use one from picker.
1546            if let Some(toolchain) = toolchain {
1547                if !user_settings.is_object() {
1548                    user_settings = Value::Object(serde_json::Map::default());
1549                }
1550                let object = user_settings.as_object_mut().unwrap();
1551                if let Some(python) = object
1552                    .entry("plugins")
1553                    .or_insert(Value::Object(serde_json::Map::default()))
1554                    .as_object_mut()
1555                {
1556                    if let Some(jedi) = python
1557                        .entry("jedi")
1558                        .or_insert(Value::Object(serde_json::Map::default()))
1559                        .as_object_mut()
1560                    {
1561                        jedi.entry("environment".to_string())
1562                            .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1563                    }
1564                    if let Some(pylint) = python
1565                        .entry("pylsp_mypy")
1566                        .or_insert(Value::Object(serde_json::Map::default()))
1567                        .as_object_mut()
1568                    {
1569                        pylint.entry("overrides".to_string()).or_insert_with(|| {
1570                            Value::Array(vec![
1571                                Value::String("--python-executable".into()),
1572                                Value::String(toolchain.path.into()),
1573                                Value::String("--cache-dir=/dev/null".into()),
1574                                Value::Bool(true),
1575                            ])
1576                        });
1577                    }
1578                }
1579            }
1580            user_settings = Value::Object(serde_json::Map::from_iter([(
1581                "pylsp".to_string(),
1582                user_settings,
1583            )]));
1584
1585            user_settings
1586        })
1587    }
1588}
1589
1590impl LspInstaller for PyLspAdapter {
1591    type BinaryVersion = ();
1592    async fn check_if_user_installed(
1593        &self,
1594        delegate: &dyn LspAdapterDelegate,
1595        toolchain: Option<Toolchain>,
1596        _: &AsyncApp,
1597    ) -> Option<LanguageServerBinary> {
1598        if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
1599            let env = delegate.shell_env().await;
1600            Some(LanguageServerBinary {
1601                path: pylsp_bin,
1602                env: Some(env),
1603                arguments: vec![],
1604            })
1605        } else {
1606            let toolchain = toolchain?;
1607            let pylsp_path = Path::new(toolchain.path.as_ref()).parent()?.join("pylsp");
1608            pylsp_path.exists().then(|| LanguageServerBinary {
1609                path: toolchain.path.to_string().into(),
1610                arguments: vec![pylsp_path.into()],
1611                env: None,
1612            })
1613        }
1614    }
1615
1616    async fn fetch_latest_server_version(
1617        &self,
1618        _: &dyn LspAdapterDelegate,
1619        _: bool,
1620        _: &mut AsyncApp,
1621    ) -> Result<()> {
1622        Ok(())
1623    }
1624
1625    async fn fetch_server_binary(
1626        &self,
1627        _: (),
1628        _: PathBuf,
1629        delegate: &dyn LspAdapterDelegate,
1630    ) -> Result<LanguageServerBinary> {
1631        let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
1632        let pip_path = venv.join(BINARY_DIR).join("pip3");
1633        ensure!(
1634            util::command::new_smol_command(pip_path.as_path())
1635                .arg("install")
1636                .arg("python-lsp-server[all]")
1637                .arg("--upgrade")
1638                .output()
1639                .await?
1640                .status
1641                .success(),
1642            "python-lsp-server[all] installation failed"
1643        );
1644        ensure!(
1645            util::command::new_smol_command(pip_path)
1646                .arg("install")
1647                .arg("pylsp-mypy")
1648                .arg("--upgrade")
1649                .output()
1650                .await?
1651                .status
1652                .success(),
1653            "pylsp-mypy installation failed"
1654        );
1655        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1656        ensure!(
1657            delegate.which(pylsp.as_os_str()).await.is_some(),
1658            "pylsp installation was incomplete"
1659        );
1660        Ok(LanguageServerBinary {
1661            path: pylsp,
1662            env: None,
1663            arguments: vec![],
1664        })
1665    }
1666
1667    async fn cached_server_binary(
1668        &self,
1669        _: PathBuf,
1670        delegate: &dyn LspAdapterDelegate,
1671    ) -> Option<LanguageServerBinary> {
1672        let venv = self.base_venv(delegate).await.ok()?;
1673        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1674        delegate.which(pylsp.as_os_str()).await?;
1675        Some(LanguageServerBinary {
1676            path: pylsp,
1677            env: None,
1678            arguments: vec![],
1679        })
1680    }
1681}
1682
1683pub(crate) struct BasedPyrightLspAdapter {
1684    node: NodeRuntime,
1685}
1686
1687impl BasedPyrightLspAdapter {
1688    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("basedpyright");
1689    const BINARY_NAME: &'static str = "basedpyright-langserver";
1690    const SERVER_PATH: &str = "node_modules/basedpyright/langserver.index.js";
1691    const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "basedpyright/langserver.index.js";
1692
1693    pub(crate) fn new(node: NodeRuntime) -> Self {
1694        BasedPyrightLspAdapter { node }
1695    }
1696
1697    async fn get_cached_server_binary(
1698        container_dir: PathBuf,
1699        node: &NodeRuntime,
1700    ) -> Option<LanguageServerBinary> {
1701        let server_path = container_dir.join(Self::SERVER_PATH);
1702        if server_path.exists() {
1703            Some(LanguageServerBinary {
1704                path: node.binary_path().await.log_err()?,
1705                env: None,
1706                arguments: vec![server_path.into(), "--stdio".into()],
1707            })
1708        } else {
1709            log::error!("missing executable in directory {:?}", server_path);
1710            None
1711        }
1712    }
1713}
1714
1715#[async_trait(?Send)]
1716impl LspAdapter for BasedPyrightLspAdapter {
1717    fn name(&self) -> LanguageServerName {
1718        Self::SERVER_NAME
1719    }
1720
1721    async fn initialization_options(
1722        self: Arc<Self>,
1723        _: &Arc<dyn LspAdapterDelegate>,
1724    ) -> Result<Option<Value>> {
1725        // Provide minimal initialization options
1726        // Virtual environment configuration will be handled through workspace configuration
1727        Ok(Some(json!({
1728            "python": {
1729                "analysis": {
1730                    "autoSearchPaths": true,
1731                    "useLibraryCodeForTypes": true,
1732                    "autoImportCompletions": true
1733                }
1734            }
1735        })))
1736    }
1737
1738    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
1739        process_pyright_completions(items);
1740    }
1741
1742    async fn label_for_completion(
1743        &self,
1744        item: &lsp::CompletionItem,
1745        language: &Arc<language::Language>,
1746    ) -> Option<language::CodeLabel> {
1747        let label = &item.label;
1748        let label_len = label.len();
1749        let grammar = language.grammar()?;
1750        let highlight_id = match item.kind? {
1751            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method"),
1752            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function"),
1753            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type"),
1754            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant"),
1755            lsp::CompletionItemKind::VARIABLE => grammar.highlight_id_for_name("variable"),
1756            _ => {
1757                return None;
1758            }
1759        };
1760        let mut text = label.clone();
1761        if let Some(completion_details) = item
1762            .label_details
1763            .as_ref()
1764            .and_then(|details| details.description.as_ref())
1765        {
1766            write!(&mut text, " {}", completion_details).ok();
1767        }
1768        Some(language::CodeLabel::filtered(
1769            text,
1770            label_len,
1771            item.filter_text.as_deref(),
1772            highlight_id
1773                .map(|id| (0..label.len(), id))
1774                .into_iter()
1775                .collect(),
1776        ))
1777    }
1778
1779    async fn label_for_symbol(
1780        &self,
1781        name: &str,
1782        kind: lsp::SymbolKind,
1783        language: &Arc<language::Language>,
1784    ) -> Option<language::CodeLabel> {
1785        let (text, filter_range, display_range) = match kind {
1786            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1787                let text = format!("def {}():\n", name);
1788                let filter_range = 4..4 + name.len();
1789                let display_range = 0..filter_range.end;
1790                (text, filter_range, display_range)
1791            }
1792            lsp::SymbolKind::CLASS => {
1793                let text = format!("class {}:", name);
1794                let filter_range = 6..6 + name.len();
1795                let display_range = 0..filter_range.end;
1796                (text, filter_range, display_range)
1797            }
1798            lsp::SymbolKind::CONSTANT => {
1799                let text = format!("{} = 0", name);
1800                let filter_range = 0..name.len();
1801                let display_range = 0..filter_range.end;
1802                (text, filter_range, display_range)
1803            }
1804            _ => return None,
1805        };
1806        Some(language::CodeLabel::new(
1807            text[display_range.clone()].to_string(),
1808            filter_range,
1809            language.highlight_text(&text.as_str().into(), display_range),
1810        ))
1811    }
1812
1813    async fn workspace_configuration(
1814        self: Arc<Self>,
1815        adapter: &Arc<dyn LspAdapterDelegate>,
1816        toolchain: Option<Toolchain>,
1817        cx: &mut AsyncApp,
1818    ) -> Result<Value> {
1819        cx.update(move |cx| {
1820            let mut user_settings =
1821                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1822                    .and_then(|s| s.settings.clone())
1823                    .unwrap_or_default();
1824
1825            // If we have a detected toolchain, configure Pyright to use it
1826            if let Some(toolchain) = toolchain
1827                && let Ok(env) = serde_json::from_value::<
1828                    pet_core::python_environment::PythonEnvironment,
1829                >(toolchain.as_json.clone())
1830            {
1831                if !user_settings.is_object() {
1832                    user_settings = Value::Object(serde_json::Map::default());
1833                }
1834                let object = user_settings.as_object_mut().unwrap();
1835
1836                let interpreter_path = toolchain.path.to_string();
1837                if let Some(venv_dir) = env.prefix {
1838                    // Set venvPath and venv at the root level
1839                    // This matches the format of a pyrightconfig.json file
1840                    if let Some(parent) = venv_dir.parent() {
1841                        // Use relative path if the venv is inside the workspace
1842                        let venv_path = if parent == adapter.worktree_root_path() {
1843                            ".".to_string()
1844                        } else {
1845                            parent.to_string_lossy().into_owned()
1846                        };
1847                        object.insert("venvPath".to_string(), Value::String(venv_path));
1848                    }
1849
1850                    if let Some(venv_name) = venv_dir.file_name() {
1851                        object.insert(
1852                            "venv".to_owned(),
1853                            Value::String(venv_name.to_string_lossy().into_owned()),
1854                        );
1855                    }
1856                }
1857
1858                // Set both pythonPath and defaultInterpreterPath for compatibility
1859                if let Some(python) = object
1860                    .entry("python")
1861                    .or_insert(Value::Object(serde_json::Map::default()))
1862                    .as_object_mut()
1863                {
1864                    python.insert(
1865                        "pythonPath".to_owned(),
1866                        Value::String(interpreter_path.clone()),
1867                    );
1868                    python.insert(
1869                        "defaultInterpreterPath".to_owned(),
1870                        Value::String(interpreter_path),
1871                    );
1872                }
1873                // Basedpyright by default uses `strict` type checking, we tone it down as to not surpris users
1874                maybe!({
1875                    let analysis = object
1876                        .entry("basedpyright.analysis")
1877                        .or_insert(Value::Object(serde_json::Map::default()));
1878                    if let serde_json::map::Entry::Vacant(v) =
1879                        analysis.as_object_mut()?.entry("typeCheckingMode")
1880                    {
1881                        v.insert(Value::String("standard".to_owned()));
1882                    }
1883                    Some(())
1884                });
1885            }
1886
1887            user_settings
1888        })
1889    }
1890}
1891
1892impl LspInstaller for BasedPyrightLspAdapter {
1893    type BinaryVersion = String;
1894
1895    async fn fetch_latest_server_version(
1896        &self,
1897        _: &dyn LspAdapterDelegate,
1898        _: bool,
1899        _: &mut AsyncApp,
1900    ) -> Result<String> {
1901        self.node
1902            .npm_package_latest_version(Self::SERVER_NAME.as_ref())
1903            .await
1904    }
1905
1906    async fn check_if_user_installed(
1907        &self,
1908        delegate: &dyn LspAdapterDelegate,
1909        _: Option<Toolchain>,
1910        _: &AsyncApp,
1911    ) -> Option<LanguageServerBinary> {
1912        if let Some(path) = delegate.which(Self::BINARY_NAME.as_ref()).await {
1913            let env = delegate.shell_env().await;
1914            Some(LanguageServerBinary {
1915                path,
1916                env: Some(env),
1917                arguments: vec!["--stdio".into()],
1918            })
1919        } else {
1920            // TODO shouldn't this be self.node.binary_path()?
1921            let node = delegate.which("node".as_ref()).await?;
1922            let (node_modules_path, _) = delegate
1923                .npm_package_installed_version(Self::SERVER_NAME.as_ref())
1924                .await
1925                .log_err()??;
1926
1927            let path = node_modules_path.join(Self::NODE_MODULE_RELATIVE_SERVER_PATH);
1928
1929            let env = delegate.shell_env().await;
1930            Some(LanguageServerBinary {
1931                path: node,
1932                env: Some(env),
1933                arguments: vec![path.into(), "--stdio".into()],
1934            })
1935        }
1936    }
1937
1938    async fn fetch_server_binary(
1939        &self,
1940        latest_version: Self::BinaryVersion,
1941        container_dir: PathBuf,
1942        delegate: &dyn LspAdapterDelegate,
1943    ) -> Result<LanguageServerBinary> {
1944        let server_path = container_dir.join(Self::SERVER_PATH);
1945
1946        self.node
1947            .npm_install_packages(
1948                &container_dir,
1949                &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
1950            )
1951            .await?;
1952
1953        let env = delegate.shell_env().await;
1954        Ok(LanguageServerBinary {
1955            path: self.node.binary_path().await?,
1956            env: Some(env),
1957            arguments: vec![server_path.into(), "--stdio".into()],
1958        })
1959    }
1960
1961    async fn check_if_version_installed(
1962        &self,
1963        version: &Self::BinaryVersion,
1964        container_dir: &PathBuf,
1965        delegate: &dyn LspAdapterDelegate,
1966    ) -> Option<LanguageServerBinary> {
1967        let server_path = container_dir.join(Self::SERVER_PATH);
1968
1969        let should_install_language_server = self
1970            .node
1971            .should_install_npm_package(
1972                Self::SERVER_NAME.as_ref(),
1973                &server_path,
1974                container_dir,
1975                VersionStrategy::Latest(version),
1976            )
1977            .await;
1978
1979        if should_install_language_server {
1980            None
1981        } else {
1982            let env = delegate.shell_env().await;
1983            Some(LanguageServerBinary {
1984                path: self.node.binary_path().await.ok()?,
1985                env: Some(env),
1986                arguments: vec![server_path.into(), "--stdio".into()],
1987            })
1988        }
1989    }
1990
1991    async fn cached_server_binary(
1992        &self,
1993        container_dir: PathBuf,
1994        delegate: &dyn LspAdapterDelegate,
1995    ) -> Option<LanguageServerBinary> {
1996        let mut binary = Self::get_cached_server_binary(container_dir, &self.node).await?;
1997        binary.env = Some(delegate.shell_env().await);
1998        Some(binary)
1999    }
2000}
2001
2002pub(crate) struct RuffLspAdapter {
2003    fs: Arc<dyn Fs>,
2004}
2005
2006#[cfg(target_os = "macos")]
2007impl RuffLspAdapter {
2008    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2009    const ARCH_SERVER_NAME: &str = "apple-darwin";
2010}
2011
2012#[cfg(target_os = "linux")]
2013impl RuffLspAdapter {
2014    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2015    const ARCH_SERVER_NAME: &str = "unknown-linux-gnu";
2016}
2017
2018#[cfg(target_os = "freebsd")]
2019impl RuffLspAdapter {
2020    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
2021    const ARCH_SERVER_NAME: &str = "unknown-freebsd";
2022}
2023
2024#[cfg(target_os = "windows")]
2025impl RuffLspAdapter {
2026    const GITHUB_ASSET_KIND: AssetKind = AssetKind::Zip;
2027    const ARCH_SERVER_NAME: &str = "pc-windows-msvc";
2028}
2029
2030impl RuffLspAdapter {
2031    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("ruff");
2032
2033    pub fn new(fs: Arc<dyn Fs>) -> RuffLspAdapter {
2034        RuffLspAdapter { fs }
2035    }
2036
2037    fn build_asset_name() -> Result<(String, String)> {
2038        let arch = match consts::ARCH {
2039            "x86" => "i686",
2040            _ => consts::ARCH,
2041        };
2042        let os = Self::ARCH_SERVER_NAME;
2043        let suffix = match consts::OS {
2044            "windows" => "zip",
2045            _ => "tar.gz",
2046        };
2047        let asset_name = format!("ruff-{arch}-{os}.{suffix}");
2048        let asset_stem = format!("ruff-{arch}-{os}");
2049        Ok((asset_stem, asset_name))
2050    }
2051}
2052
2053#[async_trait(?Send)]
2054impl LspAdapter for RuffLspAdapter {
2055    fn name(&self) -> LanguageServerName {
2056        Self::SERVER_NAME
2057    }
2058}
2059
2060impl LspInstaller for RuffLspAdapter {
2061    type BinaryVersion = GitHubLspBinaryVersion;
2062    async fn check_if_user_installed(
2063        &self,
2064        delegate: &dyn LspAdapterDelegate,
2065        toolchain: Option<Toolchain>,
2066        _: &AsyncApp,
2067    ) -> Option<LanguageServerBinary> {
2068        let ruff_in_venv = if let Some(toolchain) = toolchain
2069            && toolchain.language_name.as_ref() == "Python"
2070        {
2071            Path::new(toolchain.path.as_str())
2072                .parent()
2073                .map(|path| path.join("ruff"))
2074        } else {
2075            None
2076        };
2077
2078        for path in ruff_in_venv.into_iter().chain(["ruff".into()]) {
2079            if let Some(ruff_bin) = delegate.which(path.as_os_str()).await {
2080                let env = delegate.shell_env().await;
2081                return Some(LanguageServerBinary {
2082                    path: ruff_bin,
2083                    env: Some(env),
2084                    arguments: vec!["server".into()],
2085                });
2086            }
2087        }
2088
2089        None
2090    }
2091
2092    async fn fetch_latest_server_version(
2093        &self,
2094        delegate: &dyn LspAdapterDelegate,
2095        _: bool,
2096        _: &mut AsyncApp,
2097    ) -> Result<GitHubLspBinaryVersion> {
2098        let release =
2099            latest_github_release("astral-sh/ruff", true, false, delegate.http_client()).await?;
2100        let (_, asset_name) = Self::build_asset_name()?;
2101        let asset = release
2102            .assets
2103            .into_iter()
2104            .find(|asset| asset.name == asset_name)
2105            .with_context(|| format!("no asset found matching `{asset_name:?}`"))?;
2106        Ok(GitHubLspBinaryVersion {
2107            name: release.tag_name,
2108            url: asset.browser_download_url,
2109            digest: asset.digest,
2110        })
2111    }
2112
2113    async fn fetch_server_binary(
2114        &self,
2115        latest_version: GitHubLspBinaryVersion,
2116        container_dir: PathBuf,
2117        delegate: &dyn LspAdapterDelegate,
2118    ) -> Result<LanguageServerBinary> {
2119        let GitHubLspBinaryVersion {
2120            name,
2121            url,
2122            digest: expected_digest,
2123        } = latest_version;
2124        let destination_path = container_dir.join(format!("ruff-{name}"));
2125        let server_path = match Self::GITHUB_ASSET_KIND {
2126            AssetKind::TarGz | AssetKind::Gz => destination_path
2127                .join(Self::build_asset_name()?.0)
2128                .join("ruff"),
2129            AssetKind::Zip => destination_path.clone().join("ruff.exe"),
2130        };
2131
2132        let binary = LanguageServerBinary {
2133            path: server_path.clone(),
2134            env: None,
2135            arguments: vec!["server".into()],
2136        };
2137
2138        let metadata_path = destination_path.with_extension("metadata");
2139        let metadata = GithubBinaryMetadata::read_from_file(&metadata_path)
2140            .await
2141            .ok();
2142        if let Some(metadata) = metadata {
2143            let validity_check = async || {
2144                delegate
2145                    .try_exec(LanguageServerBinary {
2146                        path: server_path.clone(),
2147                        arguments: vec!["--version".into()],
2148                        env: None,
2149                    })
2150                    .await
2151                    .inspect_err(|err| {
2152                        log::warn!("Unable to run {server_path:?} asset, redownloading: {err}",)
2153                    })
2154            };
2155            if let (Some(actual_digest), Some(expected_digest)) =
2156                (&metadata.digest, &expected_digest)
2157            {
2158                if actual_digest == expected_digest {
2159                    if validity_check().await.is_ok() {
2160                        return Ok(binary);
2161                    }
2162                } else {
2163                    log::info!(
2164                        "SHA-256 mismatch for {destination_path:?} asset, downloading new asset. Expected: {expected_digest}, Got: {actual_digest}"
2165                    );
2166                }
2167            } else if validity_check().await.is_ok() {
2168                return Ok(binary);
2169            }
2170        }
2171
2172        download_server_binary(
2173            &*delegate.http_client(),
2174            &url,
2175            expected_digest.as_deref(),
2176            &destination_path,
2177            Self::GITHUB_ASSET_KIND,
2178        )
2179        .await?;
2180        make_file_executable(&server_path).await?;
2181        remove_matching(&container_dir, |path| path != destination_path).await;
2182        GithubBinaryMetadata::write_to_file(
2183            &GithubBinaryMetadata {
2184                metadata_version: 1,
2185                digest: expected_digest,
2186            },
2187            &metadata_path,
2188        )
2189        .await?;
2190
2191        Ok(LanguageServerBinary {
2192            path: server_path,
2193            env: None,
2194            arguments: vec!["server".into()],
2195        })
2196    }
2197
2198    async fn cached_server_binary(
2199        &self,
2200        container_dir: PathBuf,
2201        _: &dyn LspAdapterDelegate,
2202    ) -> Option<LanguageServerBinary> {
2203        maybe!(async {
2204            let mut last = None;
2205            let mut entries = self.fs.read_dir(&container_dir).await?;
2206            while let Some(entry) = entries.next().await {
2207                let path = entry?;
2208                if path.extension().is_some_and(|ext| ext == "metadata") {
2209                    continue;
2210                }
2211                last = Some(path);
2212            }
2213
2214            let path = last.context("no cached binary")?;
2215            let path = match Self::GITHUB_ASSET_KIND {
2216                AssetKind::TarGz | AssetKind::Gz => {
2217                    path.join(Self::build_asset_name()?.0).join("ruff")
2218                }
2219                AssetKind::Zip => path.join("ruff.exe"),
2220            };
2221
2222            anyhow::Ok(LanguageServerBinary {
2223                path,
2224                env: None,
2225                arguments: vec!["server".into()],
2226            })
2227        })
2228        .await
2229        .log_err()
2230    }
2231}
2232
2233#[cfg(test)]
2234mod tests {
2235    use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
2236    use language::{AutoindentMode, Buffer};
2237    use settings::SettingsStore;
2238    use std::num::NonZeroU32;
2239
2240    #[gpui::test]
2241    async fn test_python_autoindent(cx: &mut TestAppContext) {
2242        cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
2243        let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
2244        cx.update(|cx| {
2245            let test_settings = SettingsStore::test(cx);
2246            cx.set_global(test_settings);
2247            language::init(cx);
2248            cx.update_global::<SettingsStore, _>(|store, cx| {
2249                store.update_user_settings(cx, |s| {
2250                    s.project.all_languages.defaults.tab_size = NonZeroU32::new(2);
2251                });
2252            });
2253        });
2254
2255        cx.new(|cx| {
2256            let mut buffer = Buffer::local("", cx).with_language(language, cx);
2257            let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
2258                let ix = buffer.len();
2259                buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
2260            };
2261
2262            // indent after "def():"
2263            append(&mut buffer, "def a():\n", cx);
2264            assert_eq!(buffer.text(), "def a():\n  ");
2265
2266            // preserve indent after blank line
2267            append(&mut buffer, "\n  ", cx);
2268            assert_eq!(buffer.text(), "def a():\n  \n  ");
2269
2270            // indent after "if"
2271            append(&mut buffer, "if a:\n  ", cx);
2272            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    ");
2273
2274            // preserve indent after statement
2275            append(&mut buffer, "b()\n", cx);
2276            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    ");
2277
2278            // preserve indent after statement
2279            append(&mut buffer, "else", cx);
2280            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    else");
2281
2282            // dedent "else""
2283            append(&mut buffer, ":", cx);
2284            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n  else:");
2285
2286            // indent lines after else
2287            append(&mut buffer, "\n", cx);
2288            assert_eq!(
2289                buffer.text(),
2290                "def a():\n  \n  if a:\n    b()\n  else:\n    "
2291            );
2292
2293            // indent after an open paren. the closing paren is not indented
2294            // because there is another token before it on the same line.
2295            append(&mut buffer, "foo(\n1)", cx);
2296            assert_eq!(
2297                buffer.text(),
2298                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n      1)"
2299            );
2300
2301            // dedent the closing paren if it is shifted to the beginning of the line
2302            let argument_ix = buffer.text().find('1').unwrap();
2303            buffer.edit(
2304                [(argument_ix..argument_ix + 1, "")],
2305                Some(AutoindentMode::EachLine),
2306                cx,
2307            );
2308            assert_eq!(
2309                buffer.text(),
2310                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )"
2311            );
2312
2313            // preserve indent after the close paren
2314            append(&mut buffer, "\n", cx);
2315            assert_eq!(
2316                buffer.text(),
2317                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n    "
2318            );
2319
2320            // manually outdent the last line
2321            let end_whitespace_ix = buffer.len() - 4;
2322            buffer.edit(
2323                [(end_whitespace_ix..buffer.len(), "")],
2324                Some(AutoindentMode::EachLine),
2325                cx,
2326            );
2327            assert_eq!(
2328                buffer.text(),
2329                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n"
2330            );
2331
2332            // preserve the newly reduced indentation on the next newline
2333            append(&mut buffer, "\n", cx);
2334            assert_eq!(
2335                buffer.text(),
2336                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n\n"
2337            );
2338
2339            // reset to a for loop statement
2340            let statement = "for i in range(10):\n  print(i)\n";
2341            buffer.edit([(0..buffer.len(), statement)], None, cx);
2342
2343            // insert single line comment after each line
2344            let eol_ixs = statement
2345                .char_indices()
2346                .filter_map(|(ix, c)| if c == '\n' { Some(ix) } else { None })
2347                .collect::<Vec<usize>>();
2348            let editions = eol_ixs
2349                .iter()
2350                .enumerate()
2351                .map(|(i, &eol_ix)| (eol_ix..eol_ix, format!(" # comment {}", i + 1)))
2352                .collect::<Vec<(std::ops::Range<usize>, String)>>();
2353            buffer.edit(editions, Some(AutoindentMode::EachLine), cx);
2354            assert_eq!(
2355                buffer.text(),
2356                "for i in range(10): # comment 1\n  print(i) # comment 2\n"
2357            );
2358
2359            // reset to a simple if statement
2360            buffer.edit([(0..buffer.len(), "if a:\n  b(\n  )")], None, cx);
2361
2362            // dedent "else" on the line after a closing paren
2363            append(&mut buffer, "\n  else:\n", cx);
2364            assert_eq!(buffer.text(), "if a:\n  b(\n  )\nelse:\n  ");
2365
2366            buffer
2367        });
2368    }
2369}