python.rs

   1use anyhow::ensure;
   2use anyhow::{Result, anyhow};
   3use async_trait::async_trait;
   4use collections::HashMap;
   5use gpui::{App, Task};
   6use gpui::{AsyncApp, SharedString};
   7use language::LanguageToolchainStore;
   8use language::Toolchain;
   9use language::ToolchainList;
  10use language::ToolchainLister;
  11use language::language_settings::language_settings;
  12use language::{ContextProvider, LspAdapter, LspAdapterDelegate};
  13use language::{LanguageName, ManifestName, ManifestProvider, ManifestQuery};
  14use lsp::LanguageServerBinary;
  15use lsp::LanguageServerName;
  16use node_runtime::NodeRuntime;
  17use pet_core::Configuration;
  18use pet_core::os_environment::Environment;
  19use pet_core::python_environment::PythonEnvironmentKind;
  20use project::Fs;
  21use project::lsp_store::language_server_settings;
  22use serde_json::{Value, json};
  23use smol::lock::OnceCell;
  24use std::cmp::Ordering;
  25
  26use std::str::FromStr;
  27use std::sync::Mutex;
  28use std::{
  29    any::Any,
  30    borrow::Cow,
  31    ffi::OsString,
  32    fmt::Write,
  33    path::{Path, PathBuf},
  34    sync::Arc,
  35};
  36use task::{TaskTemplate, TaskTemplates, VariableName};
  37use util::ResultExt;
  38
  39pub(crate) struct PyprojectTomlManifestProvider;
  40
  41impl ManifestProvider for PyprojectTomlManifestProvider {
  42    fn name(&self) -> ManifestName {
  43        SharedString::new_static("pyproject.toml").into()
  44    }
  45
  46    fn search(
  47        &self,
  48        ManifestQuery {
  49            path,
  50            depth,
  51            delegate,
  52        }: ManifestQuery,
  53    ) -> Option<Arc<Path>> {
  54        for path in path.ancestors().take(depth) {
  55            let p = path.join("pyproject.toml");
  56            if delegate.exists(&p, Some(false)) {
  57                return Some(path.into());
  58            }
  59        }
  60
  61        None
  62    }
  63}
  64
  65const SERVER_PATH: &str = "node_modules/pyright/langserver.index.js";
  66const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "pyright/langserver.index.js";
  67
  68enum TestRunner {
  69    UNITTEST,
  70    PYTEST,
  71}
  72
  73impl FromStr for TestRunner {
  74    type Err = ();
  75
  76    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
  77        match s {
  78            "unittest" => Ok(Self::UNITTEST),
  79            "pytest" => Ok(Self::PYTEST),
  80            _ => Err(()),
  81        }
  82    }
  83}
  84
  85fn server_binary_arguments(server_path: &Path) -> Vec<OsString> {
  86    vec![server_path.into(), "--stdio".into()]
  87}
  88
  89pub struct PythonLspAdapter {
  90    node: NodeRuntime,
  91}
  92
  93impl PythonLspAdapter {
  94    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pyright");
  95
  96    pub fn new(node: NodeRuntime) -> Self {
  97        PythonLspAdapter { node }
  98    }
  99}
 100
 101#[async_trait(?Send)]
 102impl LspAdapter for PythonLspAdapter {
 103    fn name(&self) -> LanguageServerName {
 104        Self::SERVER_NAME.clone()
 105    }
 106
 107    async fn check_if_user_installed(
 108        &self,
 109        delegate: &dyn LspAdapterDelegate,
 110        _: Arc<dyn LanguageToolchainStore>,
 111        _: &AsyncApp,
 112    ) -> Option<LanguageServerBinary> {
 113        if let Some(pyright_bin) = delegate.which("pyright-langserver".as_ref()).await {
 114            let env = delegate.shell_env().await;
 115            Some(LanguageServerBinary {
 116                path: pyright_bin,
 117                env: Some(env),
 118                arguments: vec!["--stdio".into()],
 119            })
 120        } else {
 121            let node = delegate.which("node".as_ref()).await?;
 122            let (node_modules_path, _) = delegate
 123                .npm_package_installed_version(Self::SERVER_NAME.as_ref())
 124                .await
 125                .log_err()??;
 126
 127            let path = node_modules_path.join(NODE_MODULE_RELATIVE_SERVER_PATH);
 128
 129            Some(LanguageServerBinary {
 130                path: node,
 131                env: None,
 132                arguments: server_binary_arguments(&path),
 133            })
 134        }
 135    }
 136
 137    async fn fetch_latest_server_version(
 138        &self,
 139        _: &dyn LspAdapterDelegate,
 140    ) -> Result<Box<dyn 'static + Any + Send>> {
 141        Ok(Box::new(
 142            self.node
 143                .npm_package_latest_version(Self::SERVER_NAME.as_ref())
 144                .await?,
 145        ) as Box<_>)
 146    }
 147
 148    async fn fetch_server_binary(
 149        &self,
 150        latest_version: Box<dyn 'static + Send + Any>,
 151        container_dir: PathBuf,
 152        _: &dyn LspAdapterDelegate,
 153    ) -> Result<LanguageServerBinary> {
 154        let latest_version = latest_version.downcast::<String>().unwrap();
 155        let server_path = container_dir.join(SERVER_PATH);
 156
 157        self.node
 158            .npm_install_packages(
 159                &container_dir,
 160                &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
 161            )
 162            .await?;
 163
 164        Ok(LanguageServerBinary {
 165            path: self.node.binary_path().await?,
 166            env: None,
 167            arguments: server_binary_arguments(&server_path),
 168        })
 169    }
 170
 171    async fn check_if_version_installed(
 172        &self,
 173        version: &(dyn 'static + Send + Any),
 174        container_dir: &PathBuf,
 175        _: &dyn LspAdapterDelegate,
 176    ) -> Option<LanguageServerBinary> {
 177        let version = version.downcast_ref::<String>().unwrap();
 178        let server_path = container_dir.join(SERVER_PATH);
 179
 180        let should_install_language_server = self
 181            .node
 182            .should_install_npm_package(
 183                Self::SERVER_NAME.as_ref(),
 184                &server_path,
 185                &container_dir,
 186                &version,
 187            )
 188            .await;
 189
 190        if should_install_language_server {
 191            None
 192        } else {
 193            Some(LanguageServerBinary {
 194                path: self.node.binary_path().await.ok()?,
 195                env: None,
 196                arguments: server_binary_arguments(&server_path),
 197            })
 198        }
 199    }
 200
 201    async fn cached_server_binary(
 202        &self,
 203        container_dir: PathBuf,
 204        _: &dyn LspAdapterDelegate,
 205    ) -> Option<LanguageServerBinary> {
 206        get_cached_server_binary(container_dir, &self.node).await
 207    }
 208
 209    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
 210        // Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
 211        // Where `XX` is the sorting category, `YYYY` is based on most recent usage,
 212        // and `name` is the symbol name itself.
 213        //
 214        // Because the symbol name is included, there generally are not ties when
 215        // sorting by the `sortText`, so the symbol's fuzzy match score is not taken
 216        // into account. Here, we remove the symbol name from the sortText in order
 217        // to allow our own fuzzy score to be used to break ties.
 218        //
 219        // see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
 220        for item in items {
 221            let Some(sort_text) = &mut item.sort_text else {
 222                continue;
 223            };
 224            let mut parts = sort_text.split('.');
 225            let Some(first) = parts.next() else { continue };
 226            let Some(second) = parts.next() else { continue };
 227            let Some(_) = parts.next() else { continue };
 228            sort_text.replace_range(first.len() + second.len() + 1.., "");
 229        }
 230    }
 231
 232    async fn label_for_completion(
 233        &self,
 234        item: &lsp::CompletionItem,
 235        language: &Arc<language::Language>,
 236    ) -> Option<language::CodeLabel> {
 237        let label = &item.label;
 238        let grammar = language.grammar()?;
 239        let highlight_id = match item.kind? {
 240            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
 241            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
 242            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
 243            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
 244            _ => return None,
 245        };
 246        Some(language::CodeLabel {
 247            text: label.clone(),
 248            runs: vec![(0..label.len(), highlight_id)],
 249            filter_range: 0..label.len(),
 250        })
 251    }
 252
 253    async fn label_for_symbol(
 254        &self,
 255        name: &str,
 256        kind: lsp::SymbolKind,
 257        language: &Arc<language::Language>,
 258    ) -> Option<language::CodeLabel> {
 259        let (text, filter_range, display_range) = match kind {
 260            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
 261                let text = format!("def {}():\n", name);
 262                let filter_range = 4..4 + name.len();
 263                let display_range = 0..filter_range.end;
 264                (text, filter_range, display_range)
 265            }
 266            lsp::SymbolKind::CLASS => {
 267                let text = format!("class {}:", name);
 268                let filter_range = 6..6 + name.len();
 269                let display_range = 0..filter_range.end;
 270                (text, filter_range, display_range)
 271            }
 272            lsp::SymbolKind::CONSTANT => {
 273                let text = format!("{} = 0", name);
 274                let filter_range = 0..name.len();
 275                let display_range = 0..filter_range.end;
 276                (text, filter_range, display_range)
 277            }
 278            _ => return None,
 279        };
 280
 281        Some(language::CodeLabel {
 282            runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
 283            text: text[display_range].to_string(),
 284            filter_range,
 285        })
 286    }
 287
 288    async fn workspace_configuration(
 289        self: Arc<Self>,
 290        _: &dyn Fs,
 291        adapter: &Arc<dyn LspAdapterDelegate>,
 292        toolchains: Arc<dyn LanguageToolchainStore>,
 293        cx: &mut AsyncApp,
 294    ) -> Result<Value> {
 295        let toolchain = toolchains
 296            .active_toolchain(
 297                adapter.worktree_id(),
 298                Arc::from("".as_ref()),
 299                LanguageName::new("Python"),
 300                cx,
 301            )
 302            .await;
 303        cx.update(move |cx| {
 304            let mut user_settings =
 305                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
 306                    .and_then(|s| s.settings.clone())
 307                    .unwrap_or_default();
 308
 309            // If python.pythonPath is not set in user config, do so using our toolchain picker.
 310            if let Some(toolchain) = toolchain {
 311                if user_settings.is_null() {
 312                    user_settings = Value::Object(serde_json::Map::default());
 313                }
 314                let object = user_settings.as_object_mut().unwrap();
 315                if let Some(python) = object
 316                    .entry("python")
 317                    .or_insert(Value::Object(serde_json::Map::default()))
 318                    .as_object_mut()
 319                {
 320                    python
 321                        .entry("pythonPath")
 322                        .or_insert(Value::String(toolchain.path.into()));
 323                }
 324            }
 325            user_settings
 326        })
 327    }
 328    fn manifest_name(&self) -> Option<ManifestName> {
 329        Some(SharedString::new_static("pyproject.toml").into())
 330    }
 331}
 332
 333async fn get_cached_server_binary(
 334    container_dir: PathBuf,
 335    node: &NodeRuntime,
 336) -> Option<LanguageServerBinary> {
 337    let server_path = container_dir.join(SERVER_PATH);
 338    if server_path.exists() {
 339        Some(LanguageServerBinary {
 340            path: node.binary_path().await.log_err()?,
 341            env: None,
 342            arguments: server_binary_arguments(&server_path),
 343        })
 344    } else {
 345        log::error!("missing executable in directory {:?}", server_path);
 346        None
 347    }
 348}
 349
 350pub(crate) struct PythonContextProvider;
 351
 352const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
 353    VariableName::Custom(Cow::Borrowed("PYTHON_TEST_TARGET"));
 354
 355const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
 356    VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
 357
 358const PYTHON_MODULE_NAME_TASK_VARIABLE: VariableName =
 359    VariableName::Custom(Cow::Borrowed("PYTHON_MODULE_NAME"));
 360
 361impl ContextProvider for PythonContextProvider {
 362    fn build_context(
 363        &self,
 364        variables: &task::TaskVariables,
 365        location: &project::Location,
 366        _: Option<HashMap<String, String>>,
 367        toolchains: Arc<dyn LanguageToolchainStore>,
 368        cx: &mut gpui::App,
 369    ) -> Task<Result<task::TaskVariables>> {
 370        let test_target = match selected_test_runner(location.buffer.read(cx).file(), cx) {
 371            TestRunner::UNITTEST => self.build_unittest_target(variables),
 372            TestRunner::PYTEST => self.build_pytest_target(variables),
 373        };
 374
 375        let module_target = self.build_module_target(variables);
 376        let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
 377
 378        cx.spawn(async move |cx| {
 379            let active_toolchain = if let Some(worktree_id) = worktree_id {
 380                toolchains
 381                    .active_toolchain(worktree_id, Arc::from("".as_ref()), "Python".into(), cx)
 382                    .await
 383                    .map_or_else(
 384                        || "python3".to_owned(),
 385                        |toolchain| format!("\"{}\"", toolchain.path),
 386                    )
 387            } else {
 388                String::from("python3")
 389            };
 390            let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
 391
 392            Ok(task::TaskVariables::from_iter(
 393                test_target
 394                    .into_iter()
 395                    .chain(module_target.into_iter())
 396                    .chain([toolchain]),
 397            ))
 398        })
 399    }
 400
 401    fn associated_tasks(
 402        &self,
 403        file: Option<Arc<dyn language::File>>,
 404        cx: &App,
 405    ) -> Option<TaskTemplates> {
 406        let test_runner = selected_test_runner(file.as_ref(), cx);
 407
 408        let mut tasks = vec![
 409            // Execute a selection
 410            TaskTemplate {
 411                label: "execute selection".to_owned(),
 412                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 413                args: vec![
 414                    "-c".to_owned(),
 415                    VariableName::SelectedText.template_value_with_whitespace(),
 416                ],
 417                ..TaskTemplate::default()
 418            },
 419            // Execute an entire file
 420            TaskTemplate {
 421                label: format!("run '{}'", VariableName::File.template_value()),
 422                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 423                args: vec![VariableName::File.template_value_with_whitespace()],
 424                ..TaskTemplate::default()
 425            },
 426            // Execute a file as module
 427            TaskTemplate {
 428                label: format!("run module '{}'", VariableName::File.template_value()),
 429                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 430                args: vec![
 431                    "-m".to_owned(),
 432                    PYTHON_MODULE_NAME_TASK_VARIABLE.template_value(),
 433                ],
 434                tags: vec!["python-module-main-method".to_owned()],
 435                ..TaskTemplate::default()
 436            },
 437        ];
 438
 439        tasks.extend(match test_runner {
 440            TestRunner::UNITTEST => {
 441                [
 442                    // Run tests for an entire file
 443                    TaskTemplate {
 444                        label: format!("unittest '{}'", VariableName::File.template_value()),
 445                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 446                        args: vec![
 447                            "-m".to_owned(),
 448                            "unittest".to_owned(),
 449                            VariableName::File.template_value_with_whitespace(),
 450                        ],
 451                        ..TaskTemplate::default()
 452                    },
 453                    // Run test(s) for a specific target within a file
 454                    TaskTemplate {
 455                        label: "unittest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 456                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 457                        args: vec![
 458                            "-m".to_owned(),
 459                            "unittest".to_owned(),
 460                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 461                        ],
 462                        tags: vec![
 463                            "python-unittest-class".to_owned(),
 464                            "python-unittest-method".to_owned(),
 465                        ],
 466                        ..TaskTemplate::default()
 467                    },
 468                ]
 469            }
 470            TestRunner::PYTEST => {
 471                [
 472                    // Run tests for an entire file
 473                    TaskTemplate {
 474                        label: format!("pytest '{}'", VariableName::File.template_value()),
 475                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 476                        args: vec![
 477                            "-m".to_owned(),
 478                            "pytest".to_owned(),
 479                            VariableName::File.template_value_with_whitespace(),
 480                        ],
 481                        ..TaskTemplate::default()
 482                    },
 483                    // Run test(s) for a specific target within a file
 484                    TaskTemplate {
 485                        label: "pytest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 486                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
 487                        args: vec![
 488                            "-m".to_owned(),
 489                            "pytest".to_owned(),
 490                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 491                        ],
 492                        tags: vec![
 493                            "python-pytest-class".to_owned(),
 494                            "python-pytest-method".to_owned(),
 495                        ],
 496                        ..TaskTemplate::default()
 497                    },
 498                ]
 499            }
 500        });
 501
 502        Some(TaskTemplates(tasks))
 503    }
 504}
 505
 506fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &App) -> TestRunner {
 507    const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
 508    language_settings(Some(LanguageName::new("Python")), location, cx)
 509        .tasks
 510        .variables
 511        .get(TEST_RUNNER_VARIABLE)
 512        .and_then(|val| TestRunner::from_str(val).ok())
 513        .unwrap_or(TestRunner::PYTEST)
 514}
 515
 516impl PythonContextProvider {
 517    fn build_unittest_target(
 518        &self,
 519        variables: &task::TaskVariables,
 520    ) -> Option<(VariableName, String)> {
 521        let python_module_name =
 522            python_module_name_from_relative_path(variables.get(&VariableName::RelativeFile)?);
 523
 524        let unittest_class_name =
 525            variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
 526
 527        let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
 528            "_unittest_method_name",
 529        )));
 530
 531        let unittest_target_str = match (unittest_class_name, unittest_method_name) {
 532            (Some(class_name), Some(method_name)) => {
 533                format!("{python_module_name}.{class_name}.{method_name}")
 534            }
 535            (Some(class_name), None) => format!("{python_module_name}.{class_name}"),
 536            (None, None) => python_module_name,
 537            // should never happen, a TestCase class is the unit of testing
 538            (None, Some(_)) => return None,
 539        };
 540
 541        Some((
 542            PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
 543            unittest_target_str,
 544        ))
 545    }
 546
 547    fn build_pytest_target(
 548        &self,
 549        variables: &task::TaskVariables,
 550    ) -> Option<(VariableName, String)> {
 551        let file_path = variables.get(&VariableName::RelativeFile)?;
 552
 553        let pytest_class_name =
 554            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
 555
 556        let pytest_method_name =
 557            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
 558
 559        let pytest_target_str = match (pytest_class_name, pytest_method_name) {
 560            (Some(class_name), Some(method_name)) => {
 561                format!("{file_path}::{class_name}::{method_name}")
 562            }
 563            (Some(class_name), None) => {
 564                format!("{file_path}::{class_name}")
 565            }
 566            (None, Some(method_name)) => {
 567                format!("{file_path}::{method_name}")
 568            }
 569            (None, None) => file_path.to_string(),
 570        };
 571
 572        Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
 573    }
 574
 575    fn build_module_target(
 576        &self,
 577        variables: &task::TaskVariables,
 578    ) -> Result<(VariableName, String)> {
 579        let python_module_name = python_module_name_from_relative_path(
 580            variables.get(&VariableName::RelativeFile).unwrap_or(""),
 581        );
 582
 583        let module_target = (PYTHON_MODULE_NAME_TASK_VARIABLE.clone(), python_module_name);
 584
 585        Ok(module_target)
 586    }
 587}
 588
 589fn python_module_name_from_relative_path(relative_path: &str) -> String {
 590    let path_with_dots = relative_path.replace('/', ".");
 591    path_with_dots
 592        .strip_suffix(".py")
 593        .unwrap_or(&path_with_dots)
 594        .to_string()
 595}
 596
 597fn python_env_kind_display(k: &PythonEnvironmentKind) -> &'static str {
 598    match k {
 599        PythonEnvironmentKind::Conda => "Conda",
 600        PythonEnvironmentKind::Pixi => "pixi",
 601        PythonEnvironmentKind::Homebrew => "Homebrew",
 602        PythonEnvironmentKind::Pyenv => "global (Pyenv)",
 603        PythonEnvironmentKind::GlobalPaths => "global",
 604        PythonEnvironmentKind::PyenvVirtualEnv => "Pyenv",
 605        PythonEnvironmentKind::Pipenv => "Pipenv",
 606        PythonEnvironmentKind::Poetry => "Poetry",
 607        PythonEnvironmentKind::MacPythonOrg => "global (Python.org)",
 608        PythonEnvironmentKind::MacCommandLineTools => "global (Command Line Tools for Xcode)",
 609        PythonEnvironmentKind::LinuxGlobal => "global",
 610        PythonEnvironmentKind::MacXCode => "global (Xcode)",
 611        PythonEnvironmentKind::Venv => "venv",
 612        PythonEnvironmentKind::VirtualEnv => "virtualenv",
 613        PythonEnvironmentKind::VirtualEnvWrapper => "virtualenvwrapper",
 614        PythonEnvironmentKind::WindowsStore => "global (Windows Store)",
 615        PythonEnvironmentKind::WindowsRegistry => "global (Windows Registry)",
 616    }
 617}
 618
 619pub(crate) struct PythonToolchainProvider {
 620    term: SharedString,
 621}
 622
 623impl Default for PythonToolchainProvider {
 624    fn default() -> Self {
 625        Self {
 626            term: SharedString::new_static("Virtual Environment"),
 627        }
 628    }
 629}
 630
 631static ENV_PRIORITY_LIST: &'static [PythonEnvironmentKind] = &[
 632    // Prioritize non-Conda environments.
 633    PythonEnvironmentKind::Poetry,
 634    PythonEnvironmentKind::Pipenv,
 635    PythonEnvironmentKind::VirtualEnvWrapper,
 636    PythonEnvironmentKind::Venv,
 637    PythonEnvironmentKind::VirtualEnv,
 638    PythonEnvironmentKind::Pixi,
 639    PythonEnvironmentKind::Conda,
 640    PythonEnvironmentKind::Pyenv,
 641    PythonEnvironmentKind::GlobalPaths,
 642    PythonEnvironmentKind::Homebrew,
 643];
 644
 645fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
 646    if let Some(kind) = kind {
 647        ENV_PRIORITY_LIST
 648            .iter()
 649            .position(|blessed_env| blessed_env == &kind)
 650            .unwrap_or(ENV_PRIORITY_LIST.len())
 651    } else {
 652        // Unknown toolchains are less useful than non-blessed ones.
 653        ENV_PRIORITY_LIST.len() + 1
 654    }
 655}
 656
 657#[async_trait]
 658impl ToolchainLister for PythonToolchainProvider {
 659    async fn list(
 660        &self,
 661        worktree_root: PathBuf,
 662        project_env: Option<HashMap<String, String>>,
 663    ) -> ToolchainList {
 664        let env = project_env.unwrap_or_default();
 665        let environment = EnvironmentApi::from_env(&env);
 666        let locators = pet::locators::create_locators(
 667            Arc::new(pet_conda::Conda::from(&environment)),
 668            Arc::new(pet_poetry::Poetry::from(&environment)),
 669            &environment,
 670        );
 671        let mut config = Configuration::default();
 672        config.workspace_directories = Some(vec![worktree_root]);
 673        for locator in locators.iter() {
 674            locator.configure(&config);
 675        }
 676
 677        let reporter = pet_reporter::collect::create_reporter();
 678        pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
 679
 680        let mut toolchains = reporter
 681            .environments
 682            .lock()
 683            .ok()
 684            .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
 685
 686        toolchains.sort_by(|lhs, rhs| {
 687            env_priority(lhs.kind)
 688                .cmp(&env_priority(rhs.kind))
 689                .then_with(|| {
 690                    if lhs.kind == Some(PythonEnvironmentKind::Conda) {
 691                        environment
 692                            .get_env_var("CONDA_PREFIX".to_string())
 693                            .map(|conda_prefix| {
 694                                let is_match = |exe: &Option<PathBuf>| {
 695                                    exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix))
 696                                };
 697                                match (is_match(&lhs.executable), is_match(&rhs.executable)) {
 698                                    (true, false) => Ordering::Less,
 699                                    (false, true) => Ordering::Greater,
 700                                    _ => Ordering::Equal,
 701                                }
 702                            })
 703                            .unwrap_or(Ordering::Equal)
 704                    } else {
 705                        Ordering::Equal
 706                    }
 707                })
 708                .then_with(|| lhs.executable.cmp(&rhs.executable))
 709        });
 710
 711        let mut toolchains: Vec<_> = toolchains
 712            .into_iter()
 713            .filter_map(|toolchain| {
 714                let mut name = String::from("Python");
 715                if let Some(ref version) = toolchain.version {
 716                    _ = write!(name, " {version}");
 717                }
 718
 719                let name_and_kind = match (&toolchain.name, &toolchain.kind) {
 720                    (Some(name), Some(kind)) => {
 721                        Some(format!("({name}; {})", python_env_kind_display(kind)))
 722                    }
 723                    (Some(name), None) => Some(format!("({name})")),
 724                    (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
 725                    (None, None) => None,
 726                };
 727
 728                if let Some(nk) = name_and_kind {
 729                    _ = write!(name, " {nk}");
 730                }
 731
 732                Some(Toolchain {
 733                    name: name.into(),
 734                    path: toolchain.executable.as_ref()?.to_str()?.to_owned().into(),
 735                    language_name: LanguageName::new("Python"),
 736                    as_json: serde_json::to_value(toolchain).ok()?,
 737                })
 738            })
 739            .collect();
 740        toolchains.dedup();
 741        ToolchainList {
 742            toolchains,
 743            default: None,
 744            groups: Default::default(),
 745        }
 746    }
 747    fn term(&self) -> SharedString {
 748        self.term.clone()
 749    }
 750}
 751
 752pub struct EnvironmentApi<'a> {
 753    global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
 754    project_env: &'a HashMap<String, String>,
 755    pet_env: pet_core::os_environment::EnvironmentApi,
 756}
 757
 758impl<'a> EnvironmentApi<'a> {
 759    pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
 760        let paths = project_env
 761            .get("PATH")
 762            .map(|p| std::env::split_paths(p).collect())
 763            .unwrap_or_default();
 764
 765        EnvironmentApi {
 766            global_search_locations: Arc::new(Mutex::new(paths)),
 767            project_env,
 768            pet_env: pet_core::os_environment::EnvironmentApi::new(),
 769        }
 770    }
 771
 772    fn user_home(&self) -> Option<PathBuf> {
 773        self.project_env
 774            .get("HOME")
 775            .or_else(|| self.project_env.get("USERPROFILE"))
 776            .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
 777            .or_else(|| self.pet_env.get_user_home())
 778    }
 779}
 780
 781impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
 782    fn get_user_home(&self) -> Option<PathBuf> {
 783        self.user_home()
 784    }
 785
 786    fn get_root(&self) -> Option<PathBuf> {
 787        None
 788    }
 789
 790    fn get_env_var(&self, key: String) -> Option<String> {
 791        self.project_env
 792            .get(&key)
 793            .cloned()
 794            .or_else(|| self.pet_env.get_env_var(key))
 795    }
 796
 797    fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
 798        if self.global_search_locations.lock().unwrap().is_empty() {
 799            let mut paths =
 800                std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
 801                    .collect::<Vec<PathBuf>>();
 802
 803            log::trace!("Env PATH: {:?}", paths);
 804            for p in self.pet_env.get_know_global_search_locations() {
 805                if !paths.contains(&p) {
 806                    paths.push(p);
 807                }
 808            }
 809
 810            let mut paths = paths
 811                .into_iter()
 812                .filter(|p| p.exists())
 813                .collect::<Vec<PathBuf>>();
 814
 815            self.global_search_locations
 816                .lock()
 817                .unwrap()
 818                .append(&mut paths);
 819        }
 820        self.global_search_locations.lock().unwrap().clone()
 821    }
 822}
 823
 824pub(crate) struct PyLspAdapter {
 825    python_venv_base: OnceCell<Result<Arc<Path>, String>>,
 826}
 827impl PyLspAdapter {
 828    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
 829    pub(crate) fn new() -> Self {
 830        Self {
 831            python_venv_base: OnceCell::new(),
 832        }
 833    }
 834    async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
 835        let python_path = Self::find_base_python(delegate)
 836            .await
 837            .ok_or_else(|| anyhow!("Could not find Python installation for PyLSP"))?;
 838        let work_dir = delegate
 839            .language_server_download_dir(&Self::SERVER_NAME)
 840            .await
 841            .ok_or_else(|| anyhow!("Could not get working directory for PyLSP"))?;
 842        let mut path = PathBuf::from(work_dir.as_ref());
 843        path.push("pylsp-venv");
 844        if !path.exists() {
 845            util::command::new_smol_command(python_path)
 846                .arg("-m")
 847                .arg("venv")
 848                .arg("pylsp-venv")
 849                .current_dir(work_dir)
 850                .spawn()?
 851                .output()
 852                .await?;
 853        }
 854
 855        Ok(path.into())
 856    }
 857    // Find "baseline", user python version from which we'll create our own venv.
 858    async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
 859        for path in ["python3", "python"] {
 860            if let Some(path) = delegate.which(path.as_ref()).await {
 861                return Some(path);
 862            }
 863        }
 864        None
 865    }
 866
 867    async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
 868        self.python_venv_base
 869            .get_or_init(move || async move {
 870                Self::ensure_venv(delegate)
 871                    .await
 872                    .map_err(|e| format!("{e}"))
 873            })
 874            .await
 875            .clone()
 876    }
 877}
 878
 879const BINARY_DIR: &str = if cfg!(target_os = "windows") {
 880    "Scripts"
 881} else {
 882    "bin"
 883};
 884
 885#[async_trait(?Send)]
 886impl LspAdapter for PyLspAdapter {
 887    fn name(&self) -> LanguageServerName {
 888        Self::SERVER_NAME.clone()
 889    }
 890
 891    async fn check_if_user_installed(
 892        &self,
 893        delegate: &dyn LspAdapterDelegate,
 894        toolchains: Arc<dyn LanguageToolchainStore>,
 895        cx: &AsyncApp,
 896    ) -> Option<LanguageServerBinary> {
 897        if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
 898            let env = delegate.shell_env().await;
 899            Some(LanguageServerBinary {
 900                path: pylsp_bin,
 901                env: Some(env),
 902                arguments: vec![],
 903            })
 904        } else {
 905            let venv = toolchains
 906                .active_toolchain(
 907                    delegate.worktree_id(),
 908                    Arc::from("".as_ref()),
 909                    LanguageName::new("Python"),
 910                    &mut cx.clone(),
 911                )
 912                .await?;
 913            let pylsp_path = Path::new(venv.path.as_ref()).parent()?.join("pylsp");
 914            pylsp_path.exists().then(|| LanguageServerBinary {
 915                path: venv.path.to_string().into(),
 916                arguments: vec![pylsp_path.into()],
 917                env: None,
 918            })
 919        }
 920    }
 921
 922    async fn fetch_latest_server_version(
 923        &self,
 924        _: &dyn LspAdapterDelegate,
 925    ) -> Result<Box<dyn 'static + Any + Send>> {
 926        Ok(Box::new(()) as Box<_>)
 927    }
 928
 929    async fn fetch_server_binary(
 930        &self,
 931        _: Box<dyn 'static + Send + Any>,
 932        _: PathBuf,
 933        delegate: &dyn LspAdapterDelegate,
 934    ) -> Result<LanguageServerBinary> {
 935        let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
 936        let pip_path = venv.join(BINARY_DIR).join("pip3");
 937        ensure!(
 938            util::command::new_smol_command(pip_path.as_path())
 939                .arg("install")
 940                .arg("python-lsp-server")
 941                .output()
 942                .await?
 943                .status
 944                .success(),
 945            "python-lsp-server installation failed"
 946        );
 947        ensure!(
 948            util::command::new_smol_command(pip_path.as_path())
 949                .arg("install")
 950                .arg("python-lsp-server[all]")
 951                .output()
 952                .await?
 953                .status
 954                .success(),
 955            "python-lsp-server[all] installation failed"
 956        );
 957        ensure!(
 958            util::command::new_smol_command(pip_path)
 959                .arg("install")
 960                .arg("pylsp-mypy")
 961                .output()
 962                .await?
 963                .status
 964                .success(),
 965            "pylsp-mypy installation failed"
 966        );
 967        let pylsp = venv.join(BINARY_DIR).join("pylsp");
 968        Ok(LanguageServerBinary {
 969            path: pylsp,
 970            env: None,
 971            arguments: vec![],
 972        })
 973    }
 974
 975    async fn cached_server_binary(
 976        &self,
 977        _: PathBuf,
 978        delegate: &dyn LspAdapterDelegate,
 979    ) -> Option<LanguageServerBinary> {
 980        let venv = self.base_venv(delegate).await.ok()?;
 981        let pylsp = venv.join(BINARY_DIR).join("pylsp");
 982        Some(LanguageServerBinary {
 983            path: pylsp,
 984            env: None,
 985            arguments: vec![],
 986        })
 987    }
 988
 989    async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
 990
 991    async fn label_for_completion(
 992        &self,
 993        item: &lsp::CompletionItem,
 994        language: &Arc<language::Language>,
 995    ) -> Option<language::CodeLabel> {
 996        let label = &item.label;
 997        let grammar = language.grammar()?;
 998        let highlight_id = match item.kind? {
 999            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1000            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1001            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1002            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1003            _ => return None,
1004        };
1005        Some(language::CodeLabel {
1006            text: label.clone(),
1007            runs: vec![(0..label.len(), highlight_id)],
1008            filter_range: 0..label.len(),
1009        })
1010    }
1011
1012    async fn label_for_symbol(
1013        &self,
1014        name: &str,
1015        kind: lsp::SymbolKind,
1016        language: &Arc<language::Language>,
1017    ) -> Option<language::CodeLabel> {
1018        let (text, filter_range, display_range) = match kind {
1019            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1020                let text = format!("def {}():\n", name);
1021                let filter_range = 4..4 + name.len();
1022                let display_range = 0..filter_range.end;
1023                (text, filter_range, display_range)
1024            }
1025            lsp::SymbolKind::CLASS => {
1026                let text = format!("class {}:", name);
1027                let filter_range = 6..6 + name.len();
1028                let display_range = 0..filter_range.end;
1029                (text, filter_range, display_range)
1030            }
1031            lsp::SymbolKind::CONSTANT => {
1032                let text = format!("{} = 0", name);
1033                let filter_range = 0..name.len();
1034                let display_range = 0..filter_range.end;
1035                (text, filter_range, display_range)
1036            }
1037            _ => return None,
1038        };
1039
1040        Some(language::CodeLabel {
1041            runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
1042            text: text[display_range].to_string(),
1043            filter_range,
1044        })
1045    }
1046
1047    async fn workspace_configuration(
1048        self: Arc<Self>,
1049        _: &dyn Fs,
1050        adapter: &Arc<dyn LspAdapterDelegate>,
1051        toolchains: Arc<dyn LanguageToolchainStore>,
1052        cx: &mut AsyncApp,
1053    ) -> Result<Value> {
1054        let toolchain = toolchains
1055            .active_toolchain(
1056                adapter.worktree_id(),
1057                Arc::from("".as_ref()),
1058                LanguageName::new("Python"),
1059                cx,
1060            )
1061            .await;
1062        cx.update(move |cx| {
1063            let mut user_settings =
1064                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1065                    .and_then(|s| s.settings.clone())
1066                    .unwrap_or_else(|| {
1067                        json!({
1068                            "plugins": {
1069                                "pycodestyle": {"enabled": false},
1070                                "rope_autoimport": {"enabled": true, "memory": true},
1071                                "pylsp_mypy": {"enabled": false}
1072                            },
1073                            "rope": {
1074                                "ropeFolder": null
1075                            },
1076                        })
1077                    });
1078
1079            // If user did not explicitly modify their python venv, use one from picker.
1080            if let Some(toolchain) = toolchain {
1081                if user_settings.is_null() {
1082                    user_settings = Value::Object(serde_json::Map::default());
1083                }
1084                let object = user_settings.as_object_mut().unwrap();
1085                if let Some(python) = object
1086                    .entry("plugins")
1087                    .or_insert(Value::Object(serde_json::Map::default()))
1088                    .as_object_mut()
1089                {
1090                    if let Some(jedi) = python
1091                        .entry("jedi")
1092                        .or_insert(Value::Object(serde_json::Map::default()))
1093                        .as_object_mut()
1094                    {
1095                        jedi.entry("environment".to_string())
1096                            .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1097                    }
1098                    if let Some(pylint) = python
1099                        .entry("pylsp_mypy")
1100                        .or_insert(Value::Object(serde_json::Map::default()))
1101                        .as_object_mut()
1102                    {
1103                        pylint.entry("overrides".to_string()).or_insert_with(|| {
1104                            Value::Array(vec![
1105                                Value::String("--python-executable".into()),
1106                                Value::String(toolchain.path.into()),
1107                                Value::String("--cache-dir=/dev/null".into()),
1108                                Value::Bool(true),
1109                            ])
1110                        });
1111                    }
1112                }
1113            }
1114            user_settings = Value::Object(serde_json::Map::from_iter([(
1115                "pylsp".to_string(),
1116                user_settings,
1117            )]));
1118
1119            user_settings
1120        })
1121    }
1122    fn manifest_name(&self) -> Option<ManifestName> {
1123        Some(SharedString::new_static("pyproject.toml").into())
1124    }
1125}
1126
1127#[cfg(test)]
1128mod tests {
1129    use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
1130    use language::{AutoindentMode, Buffer, language_settings::AllLanguageSettings};
1131    use settings::SettingsStore;
1132    use std::num::NonZeroU32;
1133
1134    #[gpui::test]
1135    async fn test_python_autoindent(cx: &mut TestAppContext) {
1136        cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1137        let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1138        cx.update(|cx| {
1139            let test_settings = SettingsStore::test(cx);
1140            cx.set_global(test_settings);
1141            language::init(cx);
1142            cx.update_global::<SettingsStore, _>(|store, cx| {
1143                store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1144                    s.defaults.tab_size = NonZeroU32::new(2);
1145                });
1146            });
1147        });
1148
1149        cx.new(|cx| {
1150            let mut buffer = Buffer::local("", cx).with_language(language, cx);
1151            let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
1152                let ix = buffer.len();
1153                buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1154            };
1155
1156            // indent after "def():"
1157            append(&mut buffer, "def a():\n", cx);
1158            assert_eq!(buffer.text(), "def a():\n  ");
1159
1160            // preserve indent after blank line
1161            append(&mut buffer, "\n  ", cx);
1162            assert_eq!(buffer.text(), "def a():\n  \n  ");
1163
1164            // indent after "if"
1165            append(&mut buffer, "if a:\n  ", cx);
1166            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    ");
1167
1168            // preserve indent after statement
1169            append(&mut buffer, "b()\n", cx);
1170            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    ");
1171
1172            // preserve indent after statement
1173            append(&mut buffer, "else", cx);
1174            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    else");
1175
1176            // dedent "else""
1177            append(&mut buffer, ":", cx);
1178            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n  else:");
1179
1180            // indent lines after else
1181            append(&mut buffer, "\n", cx);
1182            assert_eq!(
1183                buffer.text(),
1184                "def a():\n  \n  if a:\n    b()\n  else:\n    "
1185            );
1186
1187            // indent after an open paren. the closing  paren is not indented
1188            // because there is another token before it on the same line.
1189            append(&mut buffer, "foo(\n1)", cx);
1190            assert_eq!(
1191                buffer.text(),
1192                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n      1)"
1193            );
1194
1195            // dedent the closing paren if it is shifted to the beginning of the line
1196            let argument_ix = buffer.text().find('1').unwrap();
1197            buffer.edit(
1198                [(argument_ix..argument_ix + 1, "")],
1199                Some(AutoindentMode::EachLine),
1200                cx,
1201            );
1202            assert_eq!(
1203                buffer.text(),
1204                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )"
1205            );
1206
1207            // preserve indent after the close paren
1208            append(&mut buffer, "\n", cx);
1209            assert_eq!(
1210                buffer.text(),
1211                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n    "
1212            );
1213
1214            // manually outdent the last line
1215            let end_whitespace_ix = buffer.len() - 4;
1216            buffer.edit(
1217                [(end_whitespace_ix..buffer.len(), "")],
1218                Some(AutoindentMode::EachLine),
1219                cx,
1220            );
1221            assert_eq!(
1222                buffer.text(),
1223                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n"
1224            );
1225
1226            // preserve the newly reduced indentation on the next newline
1227            append(&mut buffer, "\n", cx);
1228            assert_eq!(
1229                buffer.text(),
1230                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n\n"
1231            );
1232
1233            // reset to a simple if statement
1234            buffer.edit([(0..buffer.len(), "if a:\n  b(\n  )")], None, cx);
1235
1236            // dedent "else" on the line after a closing paren
1237            append(&mut buffer, "\n  else:\n", cx);
1238            assert_eq!(buffer.text(), "if a:\n  b(\n  )\nelse:\n  ");
1239
1240            buffer
1241        });
1242    }
1243}