python.rs

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