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