python.rs

   1use anyhow::{Context as _, ensure};
   2use anyhow::{Result, anyhow};
   3use async_trait::async_trait;
   4use collections::HashMap;
   5use dap::adapters::latest_github_release;
   6use futures::{AsyncBufReadExt, StreamExt as _};
   7use gpui::{App, Task};
   8use gpui::{AsyncApp, SharedString};
   9use http_client::github::AssetKind;
  10use http_client::github::GitHubLspBinaryVersion;
  11use language::ToolchainList;
  12use language::ToolchainLister;
  13use language::language_settings::language_settings;
  14use language::{ContextLocation, LanguageToolchainStore};
  15use language::{ContextProvider, LspAdapter, LspAdapterDelegate};
  16use language::{LanguageName, ManifestName, ManifestProvider, ManifestQuery};
  17use language::{Toolchain, ToolchainMetadata};
  18use lsp::LanguageServerBinary;
  19use lsp::LanguageServerName;
  20use node_runtime::{NodeRuntime, VersionStrategy};
  21use pet_core::Configuration;
  22use pet_core::os_environment::Environment;
  23use pet_core::python_environment::{PythonEnvironment, PythonEnvironmentKind};
  24use project::Fs;
  25use project::lsp_store::language_server_settings;
  26use serde_json::{Value, json};
  27use smol::lock::OnceCell;
  28use std::cmp::Ordering;
  29use std::env::consts;
  30use util::fs::{make_file_executable, remove_matching};
  31
  32use parking_lot::Mutex;
  33use std::str::FromStr;
  34use std::{
  35    any::Any,
  36    borrow::Cow,
  37    ffi::OsString,
  38    fmt::Write,
  39    path::{Path, PathBuf},
  40    sync::Arc,
  41};
  42use task::{ShellKind, TaskTemplate, TaskTemplates, VariableName};
  43use util::{ResultExt, maybe};
  44
  45use crate::github_download::{GithubBinaryMetadata, download_server_binary};
  46
  47pub(crate) struct PyprojectTomlManifestProvider;
  48
  49impl ManifestProvider for PyprojectTomlManifestProvider {
  50    fn name(&self) -> ManifestName {
  51        SharedString::new_static("pyproject.toml").into()
  52    }
  53
  54    fn search(
  55        &self,
  56        ManifestQuery {
  57            path,
  58            depth,
  59            delegate,
  60        }: ManifestQuery,
  61    ) -> Option<Arc<Path>> {
  62        for path in path.ancestors().take(depth) {
  63            let p = path.join("pyproject.toml");
  64            if delegate.exists(&p, Some(false)) {
  65                return Some(path.into());
  66            }
  67        }
  68
  69        None
  70    }
  71}
  72
  73const SERVER_PATH: &str = "node_modules/pyright/langserver.index.js";
  74const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "pyright/langserver.index.js";
  75
  76enum TestRunner {
  77    UNITTEST,
  78    PYTEST,
  79}
  80
  81impl FromStr for TestRunner {
  82    type Err = ();
  83
  84    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
  85        match s {
  86            "unittest" => Ok(Self::UNITTEST),
  87            "pytest" => Ok(Self::PYTEST),
  88            _ => Err(()),
  89        }
  90    }
  91}
  92
  93fn server_binary_arguments(server_path: &Path) -> Vec<OsString> {
  94    vec![server_path.into(), "--stdio".into()]
  95}
  96
  97/// Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
  98/// Where `XX` is the sorting category, `YYYY` is based on most recent usage,
  99/// and `name` is the symbol name itself.
 100///
 101/// The problem with it is that Pyright adjusts the sort text based on previous resolutions (items for which we've issued `completion/resolve` call have their sortText adjusted),
 102/// which - long story short - makes completion items list non-stable. Pyright probably relies on VSCode's implementation detail.
 103/// see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
 104fn process_pyright_completions(items: &mut [lsp::CompletionItem]) {
 105    for item in items {
 106        item.sort_text.take();
 107    }
 108}
 109pub struct PythonLspAdapter {
 110    node: NodeRuntime,
 111}
 112
 113impl PythonLspAdapter {
 114    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pyright");
 115
 116    pub fn new(node: NodeRuntime) -> Self {
 117        PythonLspAdapter { node }
 118    }
 119}
 120
 121#[async_trait(?Send)]
 122impl LspAdapter for PythonLspAdapter {
 123    fn name(&self) -> LanguageServerName {
 124        Self::SERVER_NAME
 125    }
 126
 127    async fn initialization_options(
 128        self: Arc<Self>,
 129        _: &Arc<dyn LspAdapterDelegate>,
 130    ) -> Result<Option<Value>> {
 131        // Provide minimal initialization options
 132        // Virtual environment configuration will be handled through workspace configuration
 133        Ok(Some(json!({
 134            "python": {
 135                "analysis": {
 136                    "autoSearchPaths": true,
 137                    "useLibraryCodeForTypes": true,
 138                    "autoImportCompletions": true
 139                }
 140            }
 141        })))
 142    }
 143
 144    async fn check_if_user_installed(
 145        &self,
 146        delegate: &dyn LspAdapterDelegate,
 147        _: Option<Toolchain>,
 148        _: &AsyncApp,
 149    ) -> Option<LanguageServerBinary> {
 150        if let Some(pyright_bin) = delegate.which("pyright-langserver".as_ref()).await {
 151            let env = delegate.shell_env().await;
 152            Some(LanguageServerBinary {
 153                path: pyright_bin,
 154                env: Some(env),
 155                arguments: vec!["--stdio".into()],
 156            })
 157        } else {
 158            let node = delegate.which("node".as_ref()).await?;
 159            let (node_modules_path, _) = delegate
 160                .npm_package_installed_version(Self::SERVER_NAME.as_ref())
 161                .await
 162                .log_err()??;
 163
 164            let path = node_modules_path.join(NODE_MODULE_RELATIVE_SERVER_PATH);
 165
 166            let env = delegate.shell_env().await;
 167            Some(LanguageServerBinary {
 168                path: node,
 169                env: Some(env),
 170                arguments: server_binary_arguments(&path),
 171            })
 172        }
 173    }
 174
 175    async fn fetch_latest_server_version(
 176        &self,
 177        _: &dyn LspAdapterDelegate,
 178        _: &AsyncApp,
 179    ) -> Result<Box<dyn 'static + Any + Send>> {
 180        Ok(Box::new(
 181            self.node
 182                .npm_package_latest_version(Self::SERVER_NAME.as_ref())
 183                .await?,
 184        ) as Box<_>)
 185    }
 186
 187    async fn fetch_server_binary(
 188        &self,
 189        latest_version: Box<dyn 'static + Send + Any>,
 190        container_dir: PathBuf,
 191        delegate: &dyn LspAdapterDelegate,
 192    ) -> Result<LanguageServerBinary> {
 193        let latest_version = latest_version.downcast::<String>().unwrap();
 194        let server_path = container_dir.join(SERVER_PATH);
 195
 196        self.node
 197            .npm_install_packages(
 198                &container_dir,
 199                &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
 200            )
 201            .await?;
 202
 203        let env = delegate.shell_env().await;
 204        Ok(LanguageServerBinary {
 205            path: self.node.binary_path().await?,
 206            env: Some(env),
 207            arguments: server_binary_arguments(&server_path),
 208        })
 209    }
 210
 211    async fn check_if_version_installed(
 212        &self,
 213        version: &(dyn 'static + Send + Any),
 214        container_dir: &PathBuf,
 215        delegate: &dyn LspAdapterDelegate,
 216    ) -> Option<LanguageServerBinary> {
 217        let version = version.downcast_ref::<String>().unwrap();
 218        let server_path = container_dir.join(SERVER_PATH);
 219
 220        let should_install_language_server = self
 221            .node
 222            .should_install_npm_package(
 223                Self::SERVER_NAME.as_ref(),
 224                &server_path,
 225                container_dir,
 226                VersionStrategy::Latest(version),
 227            )
 228            .await;
 229
 230        if should_install_language_server {
 231            None
 232        } else {
 233            let env = delegate.shell_env().await;
 234            Some(LanguageServerBinary {
 235                path: self.node.binary_path().await.ok()?,
 236                env: Some(env),
 237                arguments: server_binary_arguments(&server_path),
 238            })
 239        }
 240    }
 241
 242    async fn cached_server_binary(
 243        &self,
 244        container_dir: PathBuf,
 245        delegate: &dyn LspAdapterDelegate,
 246    ) -> Option<LanguageServerBinary> {
 247        let mut binary = get_cached_server_binary(container_dir, &self.node).await?;
 248        binary.env = Some(delegate.shell_env().await);
 249        Some(binary)
 250    }
 251
 252    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
 253        process_pyright_completions(items);
 254    }
 255
 256    async fn label_for_completion(
 257        &self,
 258        item: &lsp::CompletionItem,
 259        language: &Arc<language::Language>,
 260    ) -> Option<language::CodeLabel> {
 261        let label = &item.label;
 262        let grammar = language.grammar()?;
 263        let highlight_id = match item.kind? {
 264            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method"),
 265            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function"),
 266            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type"),
 267            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant"),
 268            lsp::CompletionItemKind::VARIABLE => grammar.highlight_id_for_name("variable"),
 269            _ => {
 270                return None;
 271            }
 272        };
 273        let filter_range = item
 274            .filter_text
 275            .as_deref()
 276            .and_then(|filter| label.find(filter).map(|ix| ix..ix + filter.len()))
 277            .unwrap_or(0..label.len());
 278        let mut text = label.clone();
 279        if let Some(completion_details) = item
 280            .label_details
 281            .as_ref()
 282            .and_then(|details| details.description.as_ref())
 283        {
 284            write!(&mut text, " {}", completion_details).ok();
 285        }
 286        Some(language::CodeLabel {
 287            runs: highlight_id
 288                .map(|id| (0..label.len(), id))
 289                .into_iter()
 290                .collect(),
 291            text,
 292            filter_range,
 293        })
 294    }
 295
 296    async fn label_for_symbol(
 297        &self,
 298        name: &str,
 299        kind: lsp::SymbolKind,
 300        language: &Arc<language::Language>,
 301    ) -> Option<language::CodeLabel> {
 302        let (text, filter_range, display_range) = match kind {
 303            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
 304                let text = format!("def {}():\n", name);
 305                let filter_range = 4..4 + name.len();
 306                let display_range = 0..filter_range.end;
 307                (text, filter_range, display_range)
 308            }
 309            lsp::SymbolKind::CLASS => {
 310                let text = format!("class {}:", name);
 311                let filter_range = 6..6 + name.len();
 312                let display_range = 0..filter_range.end;
 313                (text, filter_range, display_range)
 314            }
 315            lsp::SymbolKind::CONSTANT => {
 316                let text = format!("{} = 0", name);
 317                let filter_range = 0..name.len();
 318                let display_range = 0..filter_range.end;
 319                (text, filter_range, display_range)
 320            }
 321            _ => return None,
 322        };
 323
 324        Some(language::CodeLabel {
 325            runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
 326            text: text[display_range].to_string(),
 327            filter_range,
 328        })
 329    }
 330
 331    async fn workspace_configuration(
 332        self: Arc<Self>,
 333
 334        adapter: &Arc<dyn LspAdapterDelegate>,
 335        toolchain: Option<Toolchain>,
 336        cx: &mut AsyncApp,
 337    ) -> Result<Value> {
 338        cx.update(move |cx| {
 339            let mut user_settings =
 340                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
 341                    .and_then(|s| s.settings.clone())
 342                    .unwrap_or_default();
 343
 344            // If we have a detected toolchain, configure Pyright to use it
 345            if let Some(toolchain) = toolchain
 346                && let Ok(env) = serde_json::from_value::<
 347                    pet_core::python_environment::PythonEnvironment,
 348                >(toolchain.as_json.clone())
 349            {
 350                if user_settings.is_null() {
 351                    user_settings = Value::Object(serde_json::Map::default());
 352                }
 353                let object = user_settings.as_object_mut().unwrap();
 354
 355                let interpreter_path = toolchain.path.to_string();
 356                if let Some(venv_dir) = env.prefix {
 357                    // Set venvPath and venv at the root level
 358                    // This matches the format of a pyrightconfig.json file
 359                    if let Some(parent) = venv_dir.parent() {
 360                        // Use relative path if the venv is inside the workspace
 361                        let venv_path = if parent == adapter.worktree_root_path() {
 362                            ".".to_string()
 363                        } else {
 364                            parent.to_string_lossy().into_owned()
 365                        };
 366                        object.insert("venvPath".to_string(), Value::String(venv_path));
 367                    }
 368
 369                    if let Some(venv_name) = venv_dir.file_name() {
 370                        object.insert(
 371                            "venv".to_owned(),
 372                            Value::String(venv_name.to_string_lossy().into_owned()),
 373                        );
 374                    }
 375                }
 376
 377                // Always set the python interpreter path
 378                // Get or create the python section
 379                let python = object
 380                    .entry("python")
 381                    .or_insert(Value::Object(serde_json::Map::default()))
 382                    .as_object_mut()
 383                    .unwrap();
 384
 385                // Set both pythonPath and defaultInterpreterPath for compatibility
 386                python.insert(
 387                    "pythonPath".to_owned(),
 388                    Value::String(interpreter_path.clone()),
 389                );
 390                python.insert(
 391                    "defaultInterpreterPath".to_owned(),
 392                    Value::String(interpreter_path),
 393                );
 394            }
 395
 396            user_settings
 397        })
 398    }
 399}
 400
 401async fn get_cached_server_binary(
 402    container_dir: PathBuf,
 403    node: &NodeRuntime,
 404) -> Option<LanguageServerBinary> {
 405    let server_path = container_dir.join(SERVER_PATH);
 406    if server_path.exists() {
 407        Some(LanguageServerBinary {
 408            path: node.binary_path().await.log_err()?,
 409            env: None,
 410            arguments: server_binary_arguments(&server_path),
 411        })
 412    } else {
 413        log::error!("missing executable in directory {:?}", server_path);
 414        None
 415    }
 416}
 417
 418pub(crate) struct PythonContextProvider;
 419
 420const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
 421    VariableName::Custom(Cow::Borrowed("PYTHON_TEST_TARGET"));
 422
 423const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
 424    VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
 425
 426const PYTHON_MODULE_NAME_TASK_VARIABLE: VariableName =
 427    VariableName::Custom(Cow::Borrowed("PYTHON_MODULE_NAME"));
 428
 429impl ContextProvider for PythonContextProvider {
 430    fn build_context(
 431        &self,
 432        variables: &task::TaskVariables,
 433        location: ContextLocation<'_>,
 434        _: Option<HashMap<String, String>>,
 435        toolchains: Arc<dyn LanguageToolchainStore>,
 436        cx: &mut gpui::App,
 437    ) -> Task<Result<task::TaskVariables>> {
 438        let test_target =
 439            match selected_test_runner(location.file_location.buffer.read(cx).file(), cx) {
 440                TestRunner::UNITTEST => self.build_unittest_target(variables),
 441                TestRunner::PYTEST => self.build_pytest_target(variables),
 442            };
 443
 444        let module_target = self.build_module_target(variables);
 445        let location_file = location.file_location.buffer.read(cx).file().cloned();
 446        let worktree_id = location_file.as_ref().map(|f| f.worktree_id(cx));
 447
 448        cx.spawn(async move |cx| {
 449            let active_toolchain = if let Some(worktree_id) = worktree_id {
 450                let file_path = location_file
 451                    .as_ref()
 452                    .and_then(|f| f.path().parent())
 453                    .map(Arc::from)
 454                    .unwrap_or_else(|| Arc::from("".as_ref()));
 455
 456                toolchains
 457                    .active_toolchain(worktree_id, file_path, "Python".into(), cx)
 458                    .await
 459                    .map_or_else(
 460                        || String::from("python3"),
 461                        |toolchain| toolchain.path.to_string(),
 462                    )
 463            } else {
 464                String::from("python3")
 465            };
 466
 467            let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
 468
 469            Ok(task::TaskVariables::from_iter(
 470                test_target
 471                    .into_iter()
 472                    .chain(module_target.into_iter())
 473                    .chain([toolchain]),
 474            ))
 475        })
 476    }
 477
 478    fn associated_tasks(
 479        &self,
 480        file: Option<Arc<dyn language::File>>,
 481        cx: &App,
 482    ) -> Task<Option<TaskTemplates>> {
 483        let test_runner = selected_test_runner(file.as_ref(), cx);
 484
 485        let mut tasks = vec![
 486            // Execute a selection
 487            TaskTemplate {
 488                label: "execute selection".to_owned(),
 489                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 490                args: vec![
 491                    "-c".to_owned(),
 492                    VariableName::SelectedText.template_value_with_whitespace(),
 493                ],
 494                cwd: Some(VariableName::WorktreeRoot.template_value()),
 495                ..TaskTemplate::default()
 496            },
 497            // Execute an entire file
 498            TaskTemplate {
 499                label: format!("run '{}'", VariableName::File.template_value()),
 500                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 501                args: vec![VariableName::File.template_value_with_whitespace()],
 502                cwd: Some(VariableName::WorktreeRoot.template_value()),
 503                ..TaskTemplate::default()
 504            },
 505            // Execute a file as module
 506            TaskTemplate {
 507                label: format!("run module '{}'", VariableName::File.template_value()),
 508                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 509                args: vec![
 510                    "-m".to_owned(),
 511                    PYTHON_MODULE_NAME_TASK_VARIABLE.template_value_with_whitespace(),
 512                ],
 513                cwd: Some(VariableName::WorktreeRoot.template_value()),
 514                tags: vec!["python-module-main-method".to_owned()],
 515                ..TaskTemplate::default()
 516            },
 517        ];
 518
 519        tasks.extend(match test_runner {
 520            TestRunner::UNITTEST => {
 521                [
 522                    // Run tests for an entire file
 523                    TaskTemplate {
 524                        label: format!("unittest '{}'", VariableName::File.template_value()),
 525                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 526                        args: vec![
 527                            "-m".to_owned(),
 528                            "unittest".to_owned(),
 529                            VariableName::File.template_value_with_whitespace(),
 530                        ],
 531                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 532                        ..TaskTemplate::default()
 533                    },
 534                    // Run test(s) for a specific target within a file
 535                    TaskTemplate {
 536                        label: "unittest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 537                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 538                        args: vec![
 539                            "-m".to_owned(),
 540                            "unittest".to_owned(),
 541                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 542                        ],
 543                        tags: vec![
 544                            "python-unittest-class".to_owned(),
 545                            "python-unittest-method".to_owned(),
 546                        ],
 547                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 548                        ..TaskTemplate::default()
 549                    },
 550                ]
 551            }
 552            TestRunner::PYTEST => {
 553                [
 554                    // Run tests for an entire file
 555                    TaskTemplate {
 556                        label: format!("pytest '{}'", VariableName::File.template_value()),
 557                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 558                        args: vec![
 559                            "-m".to_owned(),
 560                            "pytest".to_owned(),
 561                            VariableName::File.template_value_with_whitespace(),
 562                        ],
 563                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 564                        ..TaskTemplate::default()
 565                    },
 566                    // Run test(s) for a specific target within a file
 567                    TaskTemplate {
 568                        label: "pytest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
 569                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value_with_whitespace(),
 570                        args: vec![
 571                            "-m".to_owned(),
 572                            "pytest".to_owned(),
 573                            PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
 574                        ],
 575                        cwd: Some(VariableName::WorktreeRoot.template_value()),
 576                        tags: vec![
 577                            "python-pytest-class".to_owned(),
 578                            "python-pytest-method".to_owned(),
 579                        ],
 580                        ..TaskTemplate::default()
 581                    },
 582                ]
 583            }
 584        });
 585
 586        Task::ready(Some(TaskTemplates(tasks)))
 587    }
 588}
 589
 590fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &App) -> TestRunner {
 591    const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
 592    language_settings(Some(LanguageName::new("Python")), location, cx)
 593        .tasks
 594        .variables
 595        .get(TEST_RUNNER_VARIABLE)
 596        .and_then(|val| TestRunner::from_str(val).ok())
 597        .unwrap_or(TestRunner::PYTEST)
 598}
 599
 600impl PythonContextProvider {
 601    fn build_unittest_target(
 602        &self,
 603        variables: &task::TaskVariables,
 604    ) -> Option<(VariableName, String)> {
 605        let python_module_name =
 606            python_module_name_from_relative_path(variables.get(&VariableName::RelativeFile)?);
 607
 608        let unittest_class_name =
 609            variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
 610
 611        let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
 612            "_unittest_method_name",
 613        )));
 614
 615        let unittest_target_str = match (unittest_class_name, unittest_method_name) {
 616            (Some(class_name), Some(method_name)) => {
 617                format!("{python_module_name}.{class_name}.{method_name}")
 618            }
 619            (Some(class_name), None) => format!("{python_module_name}.{class_name}"),
 620            (None, None) => python_module_name,
 621            // should never happen, a TestCase class is the unit of testing
 622            (None, Some(_)) => return None,
 623        };
 624
 625        Some((
 626            PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
 627            unittest_target_str,
 628        ))
 629    }
 630
 631    fn build_pytest_target(
 632        &self,
 633        variables: &task::TaskVariables,
 634    ) -> Option<(VariableName, String)> {
 635        let file_path = variables.get(&VariableName::RelativeFile)?;
 636
 637        let pytest_class_name =
 638            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
 639
 640        let pytest_method_name =
 641            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
 642
 643        let pytest_target_str = match (pytest_class_name, pytest_method_name) {
 644            (Some(class_name), Some(method_name)) => {
 645                format!("{file_path}::{class_name}::{method_name}")
 646            }
 647            (Some(class_name), None) => {
 648                format!("{file_path}::{class_name}")
 649            }
 650            (None, Some(method_name)) => {
 651                format!("{file_path}::{method_name}")
 652            }
 653            (None, None) => file_path.to_string(),
 654        };
 655
 656        Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
 657    }
 658
 659    fn build_module_target(
 660        &self,
 661        variables: &task::TaskVariables,
 662    ) -> Result<(VariableName, String)> {
 663        let python_module_name = python_module_name_from_relative_path(
 664            variables.get(&VariableName::RelativeFile).unwrap_or(""),
 665        );
 666
 667        let module_target = (PYTHON_MODULE_NAME_TASK_VARIABLE.clone(), python_module_name);
 668
 669        Ok(module_target)
 670    }
 671}
 672
 673fn python_module_name_from_relative_path(relative_path: &str) -> String {
 674    let path_with_dots = relative_path.replace('/', ".");
 675    path_with_dots
 676        .strip_suffix(".py")
 677        .unwrap_or(&path_with_dots)
 678        .to_string()
 679}
 680
 681fn python_env_kind_display(k: &PythonEnvironmentKind) -> &'static str {
 682    match k {
 683        PythonEnvironmentKind::Conda => "Conda",
 684        PythonEnvironmentKind::Pixi => "pixi",
 685        PythonEnvironmentKind::Homebrew => "Homebrew",
 686        PythonEnvironmentKind::Pyenv => "global (Pyenv)",
 687        PythonEnvironmentKind::GlobalPaths => "global",
 688        PythonEnvironmentKind::PyenvVirtualEnv => "Pyenv",
 689        PythonEnvironmentKind::Pipenv => "Pipenv",
 690        PythonEnvironmentKind::Poetry => "Poetry",
 691        PythonEnvironmentKind::MacPythonOrg => "global (Python.org)",
 692        PythonEnvironmentKind::MacCommandLineTools => "global (Command Line Tools for Xcode)",
 693        PythonEnvironmentKind::LinuxGlobal => "global",
 694        PythonEnvironmentKind::MacXCode => "global (Xcode)",
 695        PythonEnvironmentKind::Venv => "venv",
 696        PythonEnvironmentKind::VirtualEnv => "virtualenv",
 697        PythonEnvironmentKind::VirtualEnvWrapper => "virtualenvwrapper",
 698        PythonEnvironmentKind::WindowsStore => "global (Windows Store)",
 699        PythonEnvironmentKind::WindowsRegistry => "global (Windows Registry)",
 700    }
 701}
 702
 703pub(crate) struct PythonToolchainProvider;
 704
 705static ENV_PRIORITY_LIST: &[PythonEnvironmentKind] = &[
 706    // Prioritize non-Conda environments.
 707    PythonEnvironmentKind::Poetry,
 708    PythonEnvironmentKind::Pipenv,
 709    PythonEnvironmentKind::VirtualEnvWrapper,
 710    PythonEnvironmentKind::Venv,
 711    PythonEnvironmentKind::VirtualEnv,
 712    PythonEnvironmentKind::PyenvVirtualEnv,
 713    PythonEnvironmentKind::Pixi,
 714    PythonEnvironmentKind::Conda,
 715    PythonEnvironmentKind::Pyenv,
 716    PythonEnvironmentKind::GlobalPaths,
 717    PythonEnvironmentKind::Homebrew,
 718];
 719
 720fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
 721    if let Some(kind) = kind {
 722        ENV_PRIORITY_LIST
 723            .iter()
 724            .position(|blessed_env| blessed_env == &kind)
 725            .unwrap_or(ENV_PRIORITY_LIST.len())
 726    } else {
 727        // Unknown toolchains are less useful than non-blessed ones.
 728        ENV_PRIORITY_LIST.len() + 1
 729    }
 730}
 731
 732/// Return the name of environment declared in <worktree-root/.venv.
 733///
 734/// https://virtualfish.readthedocs.io/en/latest/plugins.html#auto-activation-auto-activation
 735async fn get_worktree_venv_declaration(worktree_root: &Path) -> Option<String> {
 736    let file = async_fs::File::open(worktree_root.join(".venv"))
 737        .await
 738        .ok()?;
 739    let mut venv_name = String::new();
 740    smol::io::BufReader::new(file)
 741        .read_line(&mut venv_name)
 742        .await
 743        .ok()?;
 744    Some(venv_name.trim().to_string())
 745}
 746
 747#[async_trait]
 748impl ToolchainLister for PythonToolchainProvider {
 749    async fn list(
 750        &self,
 751        worktree_root: PathBuf,
 752        subroot_relative_path: Arc<Path>,
 753        project_env: Option<HashMap<String, String>>,
 754    ) -> ToolchainList {
 755        let env = project_env.unwrap_or_default();
 756        let environment = EnvironmentApi::from_env(&env);
 757        let locators = pet::locators::create_locators(
 758            Arc::new(pet_conda::Conda::from(&environment)),
 759            Arc::new(pet_poetry::Poetry::from(&environment)),
 760            &environment,
 761        );
 762        let mut config = Configuration::default();
 763
 764        debug_assert!(subroot_relative_path.is_relative());
 765        // `.ancestors()` will yield at least one path, so in case of empty `subroot_relative_path`, we'll just use
 766        // worktree root as the workspace directory.
 767        config.workspace_directories = Some(
 768            subroot_relative_path
 769                .ancestors()
 770                .map(|ancestor| worktree_root.join(ancestor))
 771                .collect(),
 772        );
 773        for locator in locators.iter() {
 774            locator.configure(&config);
 775        }
 776
 777        let reporter = pet_reporter::collect::create_reporter();
 778        pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
 779
 780        let mut toolchains = reporter
 781            .environments
 782            .lock()
 783            .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
 784
 785        let wr = worktree_root;
 786        let wr_venv = get_worktree_venv_declaration(&wr).await;
 787        // Sort detected environments by:
 788        //     environment name matching activation file (<workdir>/.venv)
 789        //     environment project dir matching worktree_root
 790        //     general env priority
 791        //     environment path matching the CONDA_PREFIX env var
 792        //     executable path
 793        toolchains.sort_by(|lhs, rhs| {
 794            // Compare venv names against worktree .venv file
 795            let venv_ordering =
 796                wr_venv
 797                    .as_ref()
 798                    .map_or(Ordering::Equal, |venv| match (&lhs.name, &rhs.name) {
 799                        (Some(l), Some(r)) => (r == venv).cmp(&(l == venv)),
 800                        (Some(l), None) if l == venv => Ordering::Less,
 801                        (None, Some(r)) if r == venv => Ordering::Greater,
 802                        _ => Ordering::Equal,
 803                    });
 804
 805            // Compare project paths against worktree root
 806            let proj_ordering = || match (&lhs.project, &rhs.project) {
 807                (Some(l), Some(r)) => (r == &wr).cmp(&(l == &wr)),
 808                (Some(l), None) if l == &wr => Ordering::Less,
 809                (None, Some(r)) if r == &wr => Ordering::Greater,
 810                _ => Ordering::Equal,
 811            };
 812
 813            // Compare environment priorities
 814            let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind));
 815
 816            // Compare conda prefixes
 817            let conda_ordering = || {
 818                if lhs.kind == Some(PythonEnvironmentKind::Conda) {
 819                    environment
 820                        .get_env_var("CONDA_PREFIX".to_string())
 821                        .map(|conda_prefix| {
 822                            let is_match = |exe: &Option<PathBuf>| {
 823                                exe.as_ref().is_some_and(|e| e.starts_with(&conda_prefix))
 824                            };
 825                            match (is_match(&lhs.executable), is_match(&rhs.executable)) {
 826                                (true, false) => Ordering::Less,
 827                                (false, true) => Ordering::Greater,
 828                                _ => Ordering::Equal,
 829                            }
 830                        })
 831                        .unwrap_or(Ordering::Equal)
 832                } else {
 833                    Ordering::Equal
 834                }
 835            };
 836
 837            // Compare Python executables
 838            let exe_ordering = || lhs.executable.cmp(&rhs.executable);
 839
 840            venv_ordering
 841                .then_with(proj_ordering)
 842                .then_with(priority_ordering)
 843                .then_with(conda_ordering)
 844                .then_with(exe_ordering)
 845        });
 846
 847        let mut toolchains: Vec<_> = toolchains
 848            .into_iter()
 849            .filter_map(venv_to_toolchain)
 850            .collect();
 851        toolchains.dedup();
 852        ToolchainList {
 853            toolchains,
 854            default: None,
 855            groups: Default::default(),
 856        }
 857    }
 858    fn meta(&self) -> ToolchainMetadata {
 859        ToolchainMetadata {
 860            term: SharedString::new_static("Virtual Environment"),
 861            new_toolchain_placeholder: SharedString::new_static(
 862                "A path to the python3 executable within a virtual environment, or path to virtual environment itself",
 863            ),
 864            manifest_name: ManifestName::from(SharedString::new_static("pyproject.toml")),
 865        }
 866    }
 867
 868    async fn resolve(
 869        &self,
 870        path: PathBuf,
 871        env: Option<HashMap<String, String>>,
 872    ) -> anyhow::Result<Toolchain> {
 873        let env = env.unwrap_or_default();
 874        let environment = EnvironmentApi::from_env(&env);
 875        let locators = pet::locators::create_locators(
 876            Arc::new(pet_conda::Conda::from(&environment)),
 877            Arc::new(pet_poetry::Poetry::from(&environment)),
 878            &environment,
 879        );
 880        let toolchain = pet::resolve::resolve_environment(&path, &locators, &environment)
 881            .context("Could not find a virtual environment in provided path")?;
 882        let venv = toolchain.resolved.unwrap_or(toolchain.discovered);
 883        venv_to_toolchain(venv).context("Could not convert a venv into a toolchain")
 884    }
 885
 886    async fn activation_script(
 887        &self,
 888        toolchain: &Toolchain,
 889        shell: ShellKind,
 890        fs: &dyn Fs,
 891    ) -> Vec<String> {
 892        let Ok(toolchain) = serde_json::from_value::<pet_core::python_environment::PythonEnvironment>(
 893            toolchain.as_json.clone(),
 894        ) else {
 895            return vec![];
 896        };
 897        let mut activation_script = vec![];
 898
 899        match toolchain.kind {
 900            Some(PythonEnvironmentKind::Conda) => {
 901                if let Some(name) = &toolchain.name {
 902                    activation_script.push(format!("conda activate {name}"));
 903                } else {
 904                    activation_script.push("conda activate".to_string());
 905                }
 906            }
 907            Some(PythonEnvironmentKind::Venv | PythonEnvironmentKind::VirtualEnv) => {
 908                if let Some(prefix) = &toolchain.prefix {
 909                    let activate_keyword = match shell {
 910                        ShellKind::Cmd => ".",
 911                        ShellKind::Nushell => "overlay use",
 912                        ShellKind::Powershell => ".",
 913                        ShellKind::Fish => "source",
 914                        ShellKind::Csh => "source",
 915                        ShellKind::Posix => "source",
 916                    };
 917                    let activate_script_name = match shell {
 918                        ShellKind::Posix => "activate",
 919                        ShellKind::Csh => "activate.csh",
 920                        ShellKind::Fish => "activate.fish",
 921                        ShellKind::Nushell => "activate.nu",
 922                        ShellKind::Powershell => "activate.ps1",
 923                        ShellKind::Cmd => "activate.bat",
 924                    };
 925                    let path = prefix.join(BINARY_DIR).join(activate_script_name);
 926
 927                    if let Ok(quoted) =
 928                        shlex::try_quote(&path.to_string_lossy()).map(Cow::into_owned)
 929                        && fs.is_file(&path).await
 930                    {
 931                        activation_script.push(format!("{activate_keyword} {quoted}"));
 932                    }
 933                }
 934            }
 935            Some(PythonEnvironmentKind::Pyenv) => {
 936                let Some(manager) = toolchain.manager else {
 937                    return vec![];
 938                };
 939                let version = toolchain.version.as_deref().unwrap_or("system");
 940                let pyenv = manager.executable;
 941                let pyenv = pyenv.display();
 942                activation_script.extend(match shell {
 943                    ShellKind::Fish => Some(format!("\"{pyenv}\" shell - fish {version}")),
 944                    ShellKind::Posix => Some(format!("\"{pyenv}\" shell - sh {version}")),
 945                    ShellKind::Nushell => Some(format!("\"{pyenv}\" shell - nu {version}")),
 946                    ShellKind::Powershell => None,
 947                    ShellKind::Csh => None,
 948                    ShellKind::Cmd => None,
 949                })
 950            }
 951            _ => {}
 952        }
 953        activation_script
 954    }
 955}
 956
 957fn venv_to_toolchain(venv: PythonEnvironment) -> Option<Toolchain> {
 958    let mut name = String::from("Python");
 959    if let Some(ref version) = venv.version {
 960        _ = write!(name, " {version}");
 961    }
 962
 963    let name_and_kind = match (&venv.name, &venv.kind) {
 964        (Some(name), Some(kind)) => Some(format!("({name}; {})", python_env_kind_display(kind))),
 965        (Some(name), None) => Some(format!("({name})")),
 966        (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
 967        (None, None) => None,
 968    };
 969
 970    if let Some(nk) = name_and_kind {
 971        _ = write!(name, " {nk}");
 972    }
 973
 974    Some(Toolchain {
 975        name: name.into(),
 976        path: venv.executable.as_ref()?.to_str()?.to_owned().into(),
 977        language_name: LanguageName::new("Python"),
 978        as_json: serde_json::to_value(venv).ok()?,
 979    })
 980}
 981
 982pub struct EnvironmentApi<'a> {
 983    global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
 984    project_env: &'a HashMap<String, String>,
 985    pet_env: pet_core::os_environment::EnvironmentApi,
 986}
 987
 988impl<'a> EnvironmentApi<'a> {
 989    pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
 990        let paths = project_env
 991            .get("PATH")
 992            .map(|p| std::env::split_paths(p).collect())
 993            .unwrap_or_default();
 994
 995        EnvironmentApi {
 996            global_search_locations: Arc::new(Mutex::new(paths)),
 997            project_env,
 998            pet_env: pet_core::os_environment::EnvironmentApi::new(),
 999        }
1000    }
1001
1002    fn user_home(&self) -> Option<PathBuf> {
1003        self.project_env
1004            .get("HOME")
1005            .or_else(|| self.project_env.get("USERPROFILE"))
1006            .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
1007            .or_else(|| self.pet_env.get_user_home())
1008    }
1009}
1010
1011impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
1012    fn get_user_home(&self) -> Option<PathBuf> {
1013        self.user_home()
1014    }
1015
1016    fn get_root(&self) -> Option<PathBuf> {
1017        None
1018    }
1019
1020    fn get_env_var(&self, key: String) -> Option<String> {
1021        self.project_env
1022            .get(&key)
1023            .cloned()
1024            .or_else(|| self.pet_env.get_env_var(key))
1025    }
1026
1027    fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
1028        if self.global_search_locations.lock().is_empty() {
1029            let mut paths =
1030                std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
1031                    .collect::<Vec<PathBuf>>();
1032
1033            log::trace!("Env PATH: {:?}", paths);
1034            for p in self.pet_env.get_know_global_search_locations() {
1035                if !paths.contains(&p) {
1036                    paths.push(p);
1037                }
1038            }
1039
1040            let mut paths = paths
1041                .into_iter()
1042                .filter(|p| p.exists())
1043                .collect::<Vec<PathBuf>>();
1044
1045            self.global_search_locations.lock().append(&mut paths);
1046        }
1047        self.global_search_locations.lock().clone()
1048    }
1049}
1050
1051pub(crate) struct PyLspAdapter {
1052    python_venv_base: OnceCell<Result<Arc<Path>, String>>,
1053}
1054impl PyLspAdapter {
1055    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
1056    pub(crate) fn new() -> Self {
1057        Self {
1058            python_venv_base: OnceCell::new(),
1059        }
1060    }
1061    async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
1062        let python_path = Self::find_base_python(delegate)
1063            .await
1064            .context("Could not find Python installation for PyLSP")?;
1065        let work_dir = delegate
1066            .language_server_download_dir(&Self::SERVER_NAME)
1067            .await
1068            .context("Could not get working directory for PyLSP")?;
1069        let mut path = PathBuf::from(work_dir.as_ref());
1070        path.push("pylsp-venv");
1071        if !path.exists() {
1072            util::command::new_smol_command(python_path)
1073                .arg("-m")
1074                .arg("venv")
1075                .arg("pylsp-venv")
1076                .current_dir(work_dir)
1077                .spawn()?
1078                .output()
1079                .await?;
1080        }
1081
1082        Ok(path.into())
1083    }
1084    // Find "baseline", user python version from which we'll create our own venv.
1085    async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
1086        for path in ["python3", "python"] {
1087            if let Some(path) = delegate.which(path.as_ref()).await {
1088                return Some(path);
1089            }
1090        }
1091        None
1092    }
1093
1094    async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
1095        self.python_venv_base
1096            .get_or_init(move || async move {
1097                Self::ensure_venv(delegate)
1098                    .await
1099                    .map_err(|e| format!("{e}"))
1100            })
1101            .await
1102            .clone()
1103    }
1104}
1105
1106const BINARY_DIR: &str = if cfg!(target_os = "windows") {
1107    "Scripts"
1108} else {
1109    "bin"
1110};
1111
1112#[async_trait(?Send)]
1113impl LspAdapter for PyLspAdapter {
1114    fn name(&self) -> LanguageServerName {
1115        Self::SERVER_NAME
1116    }
1117
1118    async fn check_if_user_installed(
1119        &self,
1120        delegate: &dyn LspAdapterDelegate,
1121        toolchain: Option<Toolchain>,
1122        _: &AsyncApp,
1123    ) -> Option<LanguageServerBinary> {
1124        if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
1125            let env = delegate.shell_env().await;
1126            Some(LanguageServerBinary {
1127                path: pylsp_bin,
1128                env: Some(env),
1129                arguments: vec![],
1130            })
1131        } else {
1132            let toolchain = toolchain?;
1133            let pylsp_path = Path::new(toolchain.path.as_ref()).parent()?.join("pylsp");
1134            pylsp_path.exists().then(|| LanguageServerBinary {
1135                path: toolchain.path.to_string().into(),
1136                arguments: vec![pylsp_path.into()],
1137                env: None,
1138            })
1139        }
1140    }
1141
1142    async fn fetch_latest_server_version(
1143        &self,
1144        _: &dyn LspAdapterDelegate,
1145        _: &AsyncApp,
1146    ) -> Result<Box<dyn 'static + Any + Send>> {
1147        Ok(Box::new(()) as Box<_>)
1148    }
1149
1150    async fn fetch_server_binary(
1151        &self,
1152        _: Box<dyn 'static + Send + Any>,
1153        _: PathBuf,
1154        delegate: &dyn LspAdapterDelegate,
1155    ) -> Result<LanguageServerBinary> {
1156        let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
1157        let pip_path = venv.join(BINARY_DIR).join("pip3");
1158        ensure!(
1159            util::command::new_smol_command(pip_path.as_path())
1160                .arg("install")
1161                .arg("python-lsp-server")
1162                .arg("-U")
1163                .output()
1164                .await?
1165                .status
1166                .success(),
1167            "python-lsp-server installation failed"
1168        );
1169        ensure!(
1170            util::command::new_smol_command(pip_path.as_path())
1171                .arg("install")
1172                .arg("python-lsp-server[all]")
1173                .arg("-U")
1174                .output()
1175                .await?
1176                .status
1177                .success(),
1178            "python-lsp-server[all] installation failed"
1179        );
1180        ensure!(
1181            util::command::new_smol_command(pip_path)
1182                .arg("install")
1183                .arg("pylsp-mypy")
1184                .arg("-U")
1185                .output()
1186                .await?
1187                .status
1188                .success(),
1189            "pylsp-mypy installation failed"
1190        );
1191        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1192        Ok(LanguageServerBinary {
1193            path: pylsp,
1194            env: None,
1195            arguments: vec![],
1196        })
1197    }
1198
1199    async fn cached_server_binary(
1200        &self,
1201        _: PathBuf,
1202        delegate: &dyn LspAdapterDelegate,
1203    ) -> Option<LanguageServerBinary> {
1204        let venv = self.base_venv(delegate).await.ok()?;
1205        let pylsp = venv.join(BINARY_DIR).join("pylsp");
1206        Some(LanguageServerBinary {
1207            path: pylsp,
1208            env: None,
1209            arguments: vec![],
1210        })
1211    }
1212
1213    async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
1214
1215    async fn label_for_completion(
1216        &self,
1217        item: &lsp::CompletionItem,
1218        language: &Arc<language::Language>,
1219    ) -> Option<language::CodeLabel> {
1220        let label = &item.label;
1221        let grammar = language.grammar()?;
1222        let highlight_id = match item.kind? {
1223            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1224            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1225            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1226            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1227            _ => return None,
1228        };
1229        let filter_range = item
1230            .filter_text
1231            .as_deref()
1232            .and_then(|filter| label.find(filter).map(|ix| ix..ix + filter.len()))
1233            .unwrap_or(0..label.len());
1234        Some(language::CodeLabel {
1235            text: label.clone(),
1236            runs: vec![(0..label.len(), highlight_id)],
1237            filter_range,
1238        })
1239    }
1240
1241    async fn label_for_symbol(
1242        &self,
1243        name: &str,
1244        kind: lsp::SymbolKind,
1245        language: &Arc<language::Language>,
1246    ) -> Option<language::CodeLabel> {
1247        let (text, filter_range, display_range) = match kind {
1248            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1249                let text = format!("def {}():\n", name);
1250                let filter_range = 4..4 + name.len();
1251                let display_range = 0..filter_range.end;
1252                (text, filter_range, display_range)
1253            }
1254            lsp::SymbolKind::CLASS => {
1255                let text = format!("class {}:", name);
1256                let filter_range = 6..6 + name.len();
1257                let display_range = 0..filter_range.end;
1258                (text, filter_range, display_range)
1259            }
1260            lsp::SymbolKind::CONSTANT => {
1261                let text = format!("{} = 0", name);
1262                let filter_range = 0..name.len();
1263                let display_range = 0..filter_range.end;
1264                (text, filter_range, display_range)
1265            }
1266            _ => return None,
1267        };
1268
1269        Some(language::CodeLabel {
1270            runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
1271            text: text[display_range].to_string(),
1272            filter_range,
1273        })
1274    }
1275
1276    async fn workspace_configuration(
1277        self: Arc<Self>,
1278        adapter: &Arc<dyn LspAdapterDelegate>,
1279        toolchain: Option<Toolchain>,
1280        cx: &mut AsyncApp,
1281    ) -> Result<Value> {
1282        cx.update(move |cx| {
1283            let mut user_settings =
1284                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1285                    .and_then(|s| s.settings.clone())
1286                    .unwrap_or_else(|| {
1287                        json!({
1288                            "plugins": {
1289                                "pycodestyle": {"enabled": false},
1290                                "rope_autoimport": {"enabled": true, "memory": true},
1291                                "pylsp_mypy": {"enabled": false}
1292                            },
1293                            "rope": {
1294                                "ropeFolder": null
1295                            },
1296                        })
1297                    });
1298
1299            // If user did not explicitly modify their python venv, use one from picker.
1300            if let Some(toolchain) = toolchain {
1301                if user_settings.is_null() {
1302                    user_settings = Value::Object(serde_json::Map::default());
1303                }
1304                let object = user_settings.as_object_mut().unwrap();
1305                if let Some(python) = object
1306                    .entry("plugins")
1307                    .or_insert(Value::Object(serde_json::Map::default()))
1308                    .as_object_mut()
1309                {
1310                    if let Some(jedi) = python
1311                        .entry("jedi")
1312                        .or_insert(Value::Object(serde_json::Map::default()))
1313                        .as_object_mut()
1314                    {
1315                        jedi.entry("environment".to_string())
1316                            .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1317                    }
1318                    if let Some(pylint) = python
1319                        .entry("pylsp_mypy")
1320                        .or_insert(Value::Object(serde_json::Map::default()))
1321                        .as_object_mut()
1322                    {
1323                        pylint.entry("overrides".to_string()).or_insert_with(|| {
1324                            Value::Array(vec![
1325                                Value::String("--python-executable".into()),
1326                                Value::String(toolchain.path.into()),
1327                                Value::String("--cache-dir=/dev/null".into()),
1328                                Value::Bool(true),
1329                            ])
1330                        });
1331                    }
1332                }
1333            }
1334            user_settings = Value::Object(serde_json::Map::from_iter([(
1335                "pylsp".to_string(),
1336                user_settings,
1337            )]));
1338
1339            user_settings
1340        })
1341    }
1342}
1343
1344pub(crate) struct BasedPyrightLspAdapter {
1345    python_venv_base: OnceCell<Result<Arc<Path>, String>>,
1346}
1347
1348impl BasedPyrightLspAdapter {
1349    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("basedpyright");
1350    const BINARY_NAME: &'static str = "basedpyright-langserver";
1351
1352    pub(crate) fn new() -> Self {
1353        Self {
1354            python_venv_base: OnceCell::new(),
1355        }
1356    }
1357
1358    async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
1359        let python_path = Self::find_base_python(delegate)
1360            .await
1361            .context("Could not find Python installation for basedpyright")?;
1362        let work_dir = delegate
1363            .language_server_download_dir(&Self::SERVER_NAME)
1364            .await
1365            .context("Could not get working directory for basedpyright")?;
1366        let mut path = PathBuf::from(work_dir.as_ref());
1367        path.push("basedpyright-venv");
1368        if !path.exists() {
1369            util::command::new_smol_command(python_path)
1370                .arg("-m")
1371                .arg("venv")
1372                .arg("basedpyright-venv")
1373                .current_dir(work_dir)
1374                .spawn()?
1375                .output()
1376                .await?;
1377        }
1378
1379        Ok(path.into())
1380    }
1381
1382    // Find "baseline", user python version from which we'll create our own venv.
1383    async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
1384        for path in ["python3", "python"] {
1385            if let Some(path) = delegate.which(path.as_ref()).await {
1386                return Some(path);
1387            }
1388        }
1389        None
1390    }
1391
1392    async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
1393        self.python_venv_base
1394            .get_or_init(move || async move {
1395                Self::ensure_venv(delegate)
1396                    .await
1397                    .map_err(|e| format!("{e}"))
1398            })
1399            .await
1400            .clone()
1401    }
1402}
1403
1404#[async_trait(?Send)]
1405impl LspAdapter for BasedPyrightLspAdapter {
1406    fn name(&self) -> LanguageServerName {
1407        Self::SERVER_NAME
1408    }
1409
1410    async fn initialization_options(
1411        self: Arc<Self>,
1412        _: &Arc<dyn LspAdapterDelegate>,
1413    ) -> Result<Option<Value>> {
1414        // Provide minimal initialization options
1415        // Virtual environment configuration will be handled through workspace configuration
1416        Ok(Some(json!({
1417            "python": {
1418                "analysis": {
1419                    "autoSearchPaths": true,
1420                    "useLibraryCodeForTypes": true,
1421                    "autoImportCompletions": true
1422                }
1423            }
1424        })))
1425    }
1426
1427    async fn check_if_user_installed(
1428        &self,
1429        delegate: &dyn LspAdapterDelegate,
1430        toolchain: Option<Toolchain>,
1431        _: &AsyncApp,
1432    ) -> Option<LanguageServerBinary> {
1433        if let Some(bin) = delegate.which(Self::BINARY_NAME.as_ref()).await {
1434            let env = delegate.shell_env().await;
1435            Some(LanguageServerBinary {
1436                path: bin,
1437                env: Some(env),
1438                arguments: vec!["--stdio".into()],
1439            })
1440        } else {
1441            let path = Path::new(toolchain?.path.as_ref())
1442                .parent()?
1443                .join(Self::BINARY_NAME);
1444            path.exists().then(|| LanguageServerBinary {
1445                path,
1446                arguments: vec!["--stdio".into()],
1447                env: None,
1448            })
1449        }
1450    }
1451
1452    async fn fetch_latest_server_version(
1453        &self,
1454        _: &dyn LspAdapterDelegate,
1455        _: &AsyncApp,
1456    ) -> Result<Box<dyn 'static + Any + Send>> {
1457        Ok(Box::new(()) as Box<_>)
1458    }
1459
1460    async fn fetch_server_binary(
1461        &self,
1462        _latest_version: Box<dyn 'static + Send + Any>,
1463        _container_dir: PathBuf,
1464        delegate: &dyn LspAdapterDelegate,
1465    ) -> Result<LanguageServerBinary> {
1466        let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
1467        let pip_path = venv.join(BINARY_DIR).join("pip3");
1468        ensure!(
1469            util::command::new_smol_command(pip_path.as_path())
1470                .arg("install")
1471                .arg("basedpyright")
1472                .arg("-U")
1473                .output()
1474                .await?
1475                .status
1476                .success(),
1477            "basedpyright installation failed"
1478        );
1479        let pylsp = venv.join(BINARY_DIR).join(Self::BINARY_NAME);
1480        Ok(LanguageServerBinary {
1481            path: pylsp,
1482            env: None,
1483            arguments: vec!["--stdio".into()],
1484        })
1485    }
1486
1487    async fn cached_server_binary(
1488        &self,
1489        _container_dir: PathBuf,
1490        delegate: &dyn LspAdapterDelegate,
1491    ) -> Option<LanguageServerBinary> {
1492        let venv = self.base_venv(delegate).await.ok()?;
1493        let pylsp = venv.join(BINARY_DIR).join(Self::BINARY_NAME);
1494        Some(LanguageServerBinary {
1495            path: pylsp,
1496            env: None,
1497            arguments: vec!["--stdio".into()],
1498        })
1499    }
1500
1501    async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
1502        process_pyright_completions(items);
1503    }
1504
1505    async fn label_for_completion(
1506        &self,
1507        item: &lsp::CompletionItem,
1508        language: &Arc<language::Language>,
1509    ) -> Option<language::CodeLabel> {
1510        let label = &item.label;
1511        let grammar = language.grammar()?;
1512        let highlight_id = match item.kind? {
1513            lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method"),
1514            lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function"),
1515            lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type"),
1516            lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant"),
1517            lsp::CompletionItemKind::VARIABLE => grammar.highlight_id_for_name("variable"),
1518            _ => {
1519                return None;
1520            }
1521        };
1522        let filter_range = item
1523            .filter_text
1524            .as_deref()
1525            .and_then(|filter| label.find(filter).map(|ix| ix..ix + filter.len()))
1526            .unwrap_or(0..label.len());
1527        let mut text = label.clone();
1528        if let Some(completion_details) = item
1529            .label_details
1530            .as_ref()
1531            .and_then(|details| details.description.as_ref())
1532        {
1533            write!(&mut text, " {}", completion_details).ok();
1534        }
1535        Some(language::CodeLabel {
1536            runs: highlight_id
1537                .map(|id| (0..label.len(), id))
1538                .into_iter()
1539                .collect(),
1540            text,
1541            filter_range,
1542        })
1543    }
1544
1545    async fn label_for_symbol(
1546        &self,
1547        name: &str,
1548        kind: lsp::SymbolKind,
1549        language: &Arc<language::Language>,
1550    ) -> Option<language::CodeLabel> {
1551        let (text, filter_range, display_range) = match kind {
1552            lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1553                let text = format!("def {}():\n", name);
1554                let filter_range = 4..4 + name.len();
1555                let display_range = 0..filter_range.end;
1556                (text, filter_range, display_range)
1557            }
1558            lsp::SymbolKind::CLASS => {
1559                let text = format!("class {}:", name);
1560                let filter_range = 6..6 + name.len();
1561                let display_range = 0..filter_range.end;
1562                (text, filter_range, display_range)
1563            }
1564            lsp::SymbolKind::CONSTANT => {
1565                let text = format!("{} = 0", name);
1566                let filter_range = 0..name.len();
1567                let display_range = 0..filter_range.end;
1568                (text, filter_range, display_range)
1569            }
1570            _ => return None,
1571        };
1572
1573        Some(language::CodeLabel {
1574            runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
1575            text: text[display_range].to_string(),
1576            filter_range,
1577        })
1578    }
1579
1580    async fn workspace_configuration(
1581        self: Arc<Self>,
1582        adapter: &Arc<dyn LspAdapterDelegate>,
1583        toolchain: Option<Toolchain>,
1584        cx: &mut AsyncApp,
1585    ) -> Result<Value> {
1586        cx.update(move |cx| {
1587            let mut user_settings =
1588                language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1589                    .and_then(|s| s.settings.clone())
1590                    .unwrap_or_default();
1591
1592            // If we have a detected toolchain, configure Pyright to use it
1593            if let Some(toolchain) = toolchain
1594                && let Ok(env) = serde_json::from_value::<
1595                    pet_core::python_environment::PythonEnvironment,
1596                >(toolchain.as_json.clone())
1597            {
1598                if user_settings.is_null() {
1599                    user_settings = Value::Object(serde_json::Map::default());
1600                }
1601                let object = user_settings.as_object_mut().unwrap();
1602
1603                let interpreter_path = toolchain.path.to_string();
1604                if let Some(venv_dir) = env.prefix {
1605                    // Set venvPath and venv at the root level
1606                    // This matches the format of a pyrightconfig.json file
1607                    if let Some(parent) = venv_dir.parent() {
1608                        // Use relative path if the venv is inside the workspace
1609                        let venv_path = if parent == adapter.worktree_root_path() {
1610                            ".".to_string()
1611                        } else {
1612                            parent.to_string_lossy().into_owned()
1613                        };
1614                        object.insert("venvPath".to_string(), Value::String(venv_path));
1615                    }
1616
1617                    if let Some(venv_name) = venv_dir.file_name() {
1618                        object.insert(
1619                            "venv".to_owned(),
1620                            Value::String(venv_name.to_string_lossy().into_owned()),
1621                        );
1622                    }
1623                }
1624
1625                // Set both pythonPath and defaultInterpreterPath for compatibility
1626                if let Some(python) = object
1627                    .entry("python")
1628                    .or_insert(Value::Object(serde_json::Map::default()))
1629                    .as_object_mut()
1630                {
1631                    python.insert(
1632                        "pythonPath".to_owned(),
1633                        Value::String(interpreter_path.clone()),
1634                    );
1635                    python.insert(
1636                        "defaultInterpreterPath".to_owned(),
1637                        Value::String(interpreter_path),
1638                    );
1639                }
1640                // Basedpyright by default uses `strict` type checking, we tone it down as to not surpris users
1641                maybe!({
1642                    let basedpyright = object
1643                        .entry("basedpyright")
1644                        .or_insert(Value::Object(serde_json::Map::default()));
1645                    let analysis = basedpyright
1646                        .as_object_mut()?
1647                        .entry("analysis")
1648                        .or_insert(Value::Object(serde_json::Map::default()));
1649                    if let serde_json::map::Entry::Vacant(v) =
1650                        analysis.as_object_mut()?.entry("typeCheckingMode")
1651                    {
1652                        v.insert(Value::String("standard".to_owned()));
1653                    }
1654                    Some(())
1655                });
1656            }
1657
1658            user_settings
1659        })
1660    }
1661}
1662
1663pub(crate) struct RuffLspAdapter {
1664    fs: Arc<dyn Fs>,
1665}
1666
1667#[cfg(target_os = "macos")]
1668impl RuffLspAdapter {
1669    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
1670    const ARCH_SERVER_NAME: &str = "apple-darwin";
1671}
1672
1673#[cfg(target_os = "linux")]
1674impl RuffLspAdapter {
1675    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
1676    const ARCH_SERVER_NAME: &str = "unknown-linux-gnu";
1677}
1678
1679#[cfg(target_os = "freebsd")]
1680impl RuffLspAdapter {
1681    const GITHUB_ASSET_KIND: AssetKind = AssetKind::TarGz;
1682    const ARCH_SERVER_NAME: &str = "unknown-freebsd";
1683}
1684
1685#[cfg(target_os = "windows")]
1686impl RuffLspAdapter {
1687    const GITHUB_ASSET_KIND: AssetKind = AssetKind::Zip;
1688    const ARCH_SERVER_NAME: &str = "pc-windows-msvc";
1689}
1690
1691impl RuffLspAdapter {
1692    const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("ruff");
1693
1694    pub fn new(fs: Arc<dyn Fs>) -> RuffLspAdapter {
1695        RuffLspAdapter { fs }
1696    }
1697
1698    fn build_asset_name() -> Result<(String, String)> {
1699        let arch = match consts::ARCH {
1700            "x86" => "i686",
1701            _ => consts::ARCH,
1702        };
1703        let os = Self::ARCH_SERVER_NAME;
1704        let suffix = match consts::OS {
1705            "windows" => "zip",
1706            _ => "tar.gz",
1707        };
1708        let asset_name = format!("ruff-{arch}-{os}.{suffix}");
1709        let asset_stem = format!("ruff-{arch}-{os}");
1710        Ok((asset_stem, asset_name))
1711    }
1712}
1713
1714#[async_trait(?Send)]
1715impl LspAdapter for RuffLspAdapter {
1716    fn name(&self) -> LanguageServerName {
1717        Self::SERVER_NAME
1718    }
1719
1720    async fn check_if_user_installed(
1721        &self,
1722        delegate: &dyn LspAdapterDelegate,
1723        toolchain: Option<Toolchain>,
1724        _: &AsyncApp,
1725    ) -> Option<LanguageServerBinary> {
1726        let ruff_in_venv = if let Some(toolchain) = toolchain
1727            && toolchain.language_name.as_ref() == "Python"
1728        {
1729            Path::new(toolchain.path.as_str())
1730                .parent()
1731                .map(|path| path.join("ruff"))
1732        } else {
1733            None
1734        };
1735
1736        for path in ruff_in_venv.into_iter().chain(["ruff".into()]) {
1737            if let Some(ruff_bin) = delegate.which(path.as_os_str()).await {
1738                let env = delegate.shell_env().await;
1739                return Some(LanguageServerBinary {
1740                    path: ruff_bin,
1741                    env: Some(env),
1742                    arguments: vec!["server".into()],
1743                });
1744            }
1745        }
1746
1747        None
1748    }
1749
1750    async fn fetch_latest_server_version(
1751        &self,
1752        delegate: &dyn LspAdapterDelegate,
1753        _: &AsyncApp,
1754    ) -> Result<Box<dyn 'static + Send + Any>> {
1755        let release =
1756            latest_github_release("astral-sh/ruff", true, false, delegate.http_client()).await?;
1757        let (_, asset_name) = Self::build_asset_name()?;
1758        let asset = release
1759            .assets
1760            .into_iter()
1761            .find(|asset| asset.name == asset_name)
1762            .with_context(|| format!("no asset found matching `{asset_name:?}`"))?;
1763        Ok(Box::new(GitHubLspBinaryVersion {
1764            name: release.tag_name,
1765            url: asset.browser_download_url,
1766            digest: asset.digest,
1767        }))
1768    }
1769
1770    async fn fetch_server_binary(
1771        &self,
1772        latest_version: Box<dyn 'static + Send + Any>,
1773        container_dir: PathBuf,
1774        delegate: &dyn LspAdapterDelegate,
1775    ) -> Result<LanguageServerBinary> {
1776        let GitHubLspBinaryVersion {
1777            name,
1778            url,
1779            digest: expected_digest,
1780        } = *latest_version.downcast::<GitHubLspBinaryVersion>().unwrap();
1781        let destination_path = container_dir.join(format!("ruff-{name}"));
1782        let server_path = match Self::GITHUB_ASSET_KIND {
1783            AssetKind::TarGz | AssetKind::Gz => destination_path
1784                .join(Self::build_asset_name()?.0)
1785                .join("ruff"),
1786            AssetKind::Zip => destination_path.clone().join("ruff.exe"),
1787        };
1788
1789        let binary = LanguageServerBinary {
1790            path: server_path.clone(),
1791            env: None,
1792            arguments: vec!["server".into()],
1793        };
1794
1795        let metadata_path = destination_path.with_extension("metadata");
1796        let metadata = GithubBinaryMetadata::read_from_file(&metadata_path)
1797            .await
1798            .ok();
1799        if let Some(metadata) = metadata {
1800            let validity_check = async || {
1801                delegate
1802                    .try_exec(LanguageServerBinary {
1803                        path: server_path.clone(),
1804                        arguments: vec!["--version".into()],
1805                        env: None,
1806                    })
1807                    .await
1808                    .inspect_err(|err| {
1809                        log::warn!("Unable to run {server_path:?} asset, redownloading: {err}",)
1810                    })
1811            };
1812            if let (Some(actual_digest), Some(expected_digest)) =
1813                (&metadata.digest, &expected_digest)
1814            {
1815                if actual_digest == expected_digest {
1816                    if validity_check().await.is_ok() {
1817                        return Ok(binary);
1818                    }
1819                } else {
1820                    log::info!(
1821                        "SHA-256 mismatch for {destination_path:?} asset, downloading new asset. Expected: {expected_digest}, Got: {actual_digest}"
1822                    );
1823                }
1824            } else if validity_check().await.is_ok() {
1825                return Ok(binary);
1826            }
1827        }
1828
1829        download_server_binary(
1830            delegate,
1831            &url,
1832            expected_digest.as_deref(),
1833            &destination_path,
1834            Self::GITHUB_ASSET_KIND,
1835        )
1836        .await?;
1837        make_file_executable(&server_path).await?;
1838        remove_matching(&container_dir, |path| path != destination_path).await;
1839        GithubBinaryMetadata::write_to_file(
1840            &GithubBinaryMetadata {
1841                metadata_version: 1,
1842                digest: expected_digest,
1843            },
1844            &metadata_path,
1845        )
1846        .await?;
1847
1848        Ok(LanguageServerBinary {
1849            path: server_path,
1850            env: None,
1851            arguments: vec!["server".into()],
1852        })
1853    }
1854
1855    async fn cached_server_binary(
1856        &self,
1857        container_dir: PathBuf,
1858        _: &dyn LspAdapterDelegate,
1859    ) -> Option<LanguageServerBinary> {
1860        maybe!(async {
1861            let mut last = None;
1862            let mut entries = self.fs.read_dir(&container_dir).await?;
1863            while let Some(entry) = entries.next().await {
1864                let path = entry?;
1865                if path.extension().is_some_and(|ext| ext == "metadata") {
1866                    continue;
1867                }
1868                last = Some(path);
1869            }
1870
1871            let path = last.context("no cached binary")?;
1872            let path = match Self::GITHUB_ASSET_KIND {
1873                AssetKind::TarGz | AssetKind::Gz => {
1874                    path.join(Self::build_asset_name()?.0).join("ruff")
1875                }
1876                AssetKind::Zip => path.join("ruff.exe"),
1877            };
1878
1879            anyhow::Ok(LanguageServerBinary {
1880                path,
1881                env: None,
1882                arguments: vec!["server".into()],
1883            })
1884        })
1885        .await
1886        .log_err()
1887    }
1888}
1889
1890#[cfg(test)]
1891mod tests {
1892    use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
1893    use language::{AutoindentMode, Buffer, language_settings::AllLanguageSettings};
1894    use settings::SettingsStore;
1895    use std::num::NonZeroU32;
1896
1897    #[gpui::test]
1898    async fn test_python_autoindent(cx: &mut TestAppContext) {
1899        cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1900        let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1901        cx.update(|cx| {
1902            let test_settings = SettingsStore::test(cx);
1903            cx.set_global(test_settings);
1904            language::init(cx);
1905            cx.update_global::<SettingsStore, _>(|store, cx| {
1906                store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1907                    s.defaults.tab_size = NonZeroU32::new(2);
1908                });
1909            });
1910        });
1911
1912        cx.new(|cx| {
1913            let mut buffer = Buffer::local("", cx).with_language(language, cx);
1914            let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
1915                let ix = buffer.len();
1916                buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1917            };
1918
1919            // indent after "def():"
1920            append(&mut buffer, "def a():\n", cx);
1921            assert_eq!(buffer.text(), "def a():\n  ");
1922
1923            // preserve indent after blank line
1924            append(&mut buffer, "\n  ", cx);
1925            assert_eq!(buffer.text(), "def a():\n  \n  ");
1926
1927            // indent after "if"
1928            append(&mut buffer, "if a:\n  ", cx);
1929            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    ");
1930
1931            // preserve indent after statement
1932            append(&mut buffer, "b()\n", cx);
1933            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    ");
1934
1935            // preserve indent after statement
1936            append(&mut buffer, "else", cx);
1937            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n    else");
1938
1939            // dedent "else""
1940            append(&mut buffer, ":", cx);
1941            assert_eq!(buffer.text(), "def a():\n  \n  if a:\n    b()\n  else:");
1942
1943            // indent lines after else
1944            append(&mut buffer, "\n", cx);
1945            assert_eq!(
1946                buffer.text(),
1947                "def a():\n  \n  if a:\n    b()\n  else:\n    "
1948            );
1949
1950            // indent after an open paren. the closing paren is not indented
1951            // because there is another token before it on the same line.
1952            append(&mut buffer, "foo(\n1)", cx);
1953            assert_eq!(
1954                buffer.text(),
1955                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n      1)"
1956            );
1957
1958            // dedent the closing paren if it is shifted to the beginning of the line
1959            let argument_ix = buffer.text().find('1').unwrap();
1960            buffer.edit(
1961                [(argument_ix..argument_ix + 1, "")],
1962                Some(AutoindentMode::EachLine),
1963                cx,
1964            );
1965            assert_eq!(
1966                buffer.text(),
1967                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )"
1968            );
1969
1970            // preserve indent after the close paren
1971            append(&mut buffer, "\n", cx);
1972            assert_eq!(
1973                buffer.text(),
1974                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n    "
1975            );
1976
1977            // manually outdent the last line
1978            let end_whitespace_ix = buffer.len() - 4;
1979            buffer.edit(
1980                [(end_whitespace_ix..buffer.len(), "")],
1981                Some(AutoindentMode::EachLine),
1982                cx,
1983            );
1984            assert_eq!(
1985                buffer.text(),
1986                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n"
1987            );
1988
1989            // preserve the newly reduced indentation on the next newline
1990            append(&mut buffer, "\n", cx);
1991            assert_eq!(
1992                buffer.text(),
1993                "def a():\n  \n  if a:\n    b()\n  else:\n    foo(\n    )\n\n"
1994            );
1995
1996            // reset to a for loop statement
1997            let statement = "for i in range(10):\n  print(i)\n";
1998            buffer.edit([(0..buffer.len(), statement)], None, cx);
1999
2000            // insert single line comment after each line
2001            let eol_ixs = statement
2002                .char_indices()
2003                .filter_map(|(ix, c)| if c == '\n' { Some(ix) } else { None })
2004                .collect::<Vec<usize>>();
2005            let editions = eol_ixs
2006                .iter()
2007                .enumerate()
2008                .map(|(i, &eol_ix)| (eol_ix..eol_ix, format!(" # comment {}", i + 1)))
2009                .collect::<Vec<(std::ops::Range<usize>, String)>>();
2010            buffer.edit(editions, Some(AutoindentMode::EachLine), cx);
2011            assert_eq!(
2012                buffer.text(),
2013                "for i in range(10): # comment 1\n  print(i) # comment 2\n"
2014            );
2015
2016            // reset to a simple if statement
2017            buffer.edit([(0..buffer.len(), "if a:\n  b(\n  )")], None, cx);
2018
2019            // dedent "else" on the line after a closing paren
2020            append(&mut buffer, "\n  else:\n", cx);
2021            assert_eq!(buffer.text(), "if a:\n  b(\n  )\nelse:\n  ");
2022
2023            buffer
2024        });
2025    }
2026}