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