python.rs

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