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 fn debug_adapter(&self) -> Option<String> {
479 Some("Debugpy".into())
480 }
481}
482
483fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &App) -> TestRunner {
484 const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
485 language_settings(Some(LanguageName::new("Python")), location, cx)
486 .tasks
487 .variables
488 .get(TEST_RUNNER_VARIABLE)
489 .and_then(|val| TestRunner::from_str(val).ok())
490 .unwrap_or(TestRunner::PYTEST)
491}
492
493impl PythonContextProvider {
494 fn build_unittest_target(
495 &self,
496 variables: &task::TaskVariables,
497 ) -> Option<(VariableName, String)> {
498 let python_module_name =
499 python_module_name_from_relative_path(variables.get(&VariableName::RelativeFile)?);
500
501 let unittest_class_name =
502 variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
503
504 let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
505 "_unittest_method_name",
506 )));
507
508 let unittest_target_str = match (unittest_class_name, unittest_method_name) {
509 (Some(class_name), Some(method_name)) => {
510 format!("{python_module_name}.{class_name}.{method_name}")
511 }
512 (Some(class_name), None) => format!("{python_module_name}.{class_name}"),
513 (None, None) => python_module_name,
514 // should never happen, a TestCase class is the unit of testing
515 (None, Some(_)) => return None,
516 };
517
518 Some((
519 PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
520 unittest_target_str,
521 ))
522 }
523
524 fn build_pytest_target(
525 &self,
526 variables: &task::TaskVariables,
527 ) -> Option<(VariableName, String)> {
528 let file_path = variables.get(&VariableName::RelativeFile)?;
529
530 let pytest_class_name =
531 variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
532
533 let pytest_method_name =
534 variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
535
536 let pytest_target_str = match (pytest_class_name, pytest_method_name) {
537 (Some(class_name), Some(method_name)) => {
538 format!("{file_path}::{class_name}::{method_name}")
539 }
540 (Some(class_name), None) => {
541 format!("{file_path}::{class_name}")
542 }
543 (None, Some(method_name)) => {
544 format!("{file_path}::{method_name}")
545 }
546 (None, None) => file_path.to_string(),
547 };
548
549 Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
550 }
551
552 fn build_module_target(
553 &self,
554 variables: &task::TaskVariables,
555 ) -> Result<(VariableName, String)> {
556 let python_module_name = python_module_name_from_relative_path(
557 variables.get(&VariableName::RelativeFile).unwrap_or(""),
558 );
559
560 let module_target = (PYTHON_MODULE_NAME_TASK_VARIABLE.clone(), python_module_name);
561
562 Ok(module_target)
563 }
564}
565
566fn python_module_name_from_relative_path(relative_path: &str) -> String {
567 let path_with_dots = relative_path.replace('/', ".");
568 path_with_dots
569 .strip_suffix(".py")
570 .unwrap_or(&path_with_dots)
571 .to_string()
572}
573
574fn python_env_kind_display(k: &PythonEnvironmentKind) -> &'static str {
575 match k {
576 PythonEnvironmentKind::Conda => "Conda",
577 PythonEnvironmentKind::Pixi => "pixi",
578 PythonEnvironmentKind::Homebrew => "Homebrew",
579 PythonEnvironmentKind::Pyenv => "global (Pyenv)",
580 PythonEnvironmentKind::GlobalPaths => "global",
581 PythonEnvironmentKind::PyenvVirtualEnv => "Pyenv",
582 PythonEnvironmentKind::Pipenv => "Pipenv",
583 PythonEnvironmentKind::Poetry => "Poetry",
584 PythonEnvironmentKind::MacPythonOrg => "global (Python.org)",
585 PythonEnvironmentKind::MacCommandLineTools => "global (Command Line Tools for Xcode)",
586 PythonEnvironmentKind::LinuxGlobal => "global",
587 PythonEnvironmentKind::MacXCode => "global (Xcode)",
588 PythonEnvironmentKind::Venv => "venv",
589 PythonEnvironmentKind::VirtualEnv => "virtualenv",
590 PythonEnvironmentKind::VirtualEnvWrapper => "virtualenvwrapper",
591 PythonEnvironmentKind::WindowsStore => "global (Windows Store)",
592 PythonEnvironmentKind::WindowsRegistry => "global (Windows Registry)",
593 }
594}
595
596pub(crate) struct PythonToolchainProvider {
597 term: SharedString,
598}
599
600impl Default for PythonToolchainProvider {
601 fn default() -> Self {
602 Self {
603 term: SharedString::new_static("Virtual Environment"),
604 }
605 }
606}
607
608static ENV_PRIORITY_LIST: &'static [PythonEnvironmentKind] = &[
609 // Prioritize non-Conda environments.
610 PythonEnvironmentKind::Poetry,
611 PythonEnvironmentKind::Pipenv,
612 PythonEnvironmentKind::VirtualEnvWrapper,
613 PythonEnvironmentKind::Venv,
614 PythonEnvironmentKind::VirtualEnv,
615 PythonEnvironmentKind::PyenvVirtualEnv,
616 PythonEnvironmentKind::Pixi,
617 PythonEnvironmentKind::Conda,
618 PythonEnvironmentKind::Pyenv,
619 PythonEnvironmentKind::GlobalPaths,
620 PythonEnvironmentKind::Homebrew,
621];
622
623fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
624 if let Some(kind) = kind {
625 ENV_PRIORITY_LIST
626 .iter()
627 .position(|blessed_env| blessed_env == &kind)
628 .unwrap_or(ENV_PRIORITY_LIST.len())
629 } else {
630 // Unknown toolchains are less useful than non-blessed ones.
631 ENV_PRIORITY_LIST.len() + 1
632 }
633}
634
635/// Return the name of environment declared in <worktree-root/.venv.
636///
637/// https://virtualfish.readthedocs.io/en/latest/plugins.html#auto-activation-auto-activation
638fn get_worktree_venv_declaration(worktree_root: &Path) -> Option<String> {
639 fs::File::open(worktree_root.join(".venv"))
640 .and_then(|file| {
641 let mut venv_name = String::new();
642 io::BufReader::new(file).read_line(&mut venv_name)?;
643 Ok(venv_name.trim().to_string())
644 })
645 .ok()
646}
647
648#[async_trait]
649impl ToolchainLister for PythonToolchainProvider {
650 async fn list(
651 &self,
652 worktree_root: PathBuf,
653 project_env: Option<HashMap<String, String>>,
654 ) -> ToolchainList {
655 let env = project_env.unwrap_or_default();
656 let environment = EnvironmentApi::from_env(&env);
657 let locators = pet::locators::create_locators(
658 Arc::new(pet_conda::Conda::from(&environment)),
659 Arc::new(pet_poetry::Poetry::from(&environment)),
660 &environment,
661 );
662 let mut config = Configuration::default();
663 config.workspace_directories = Some(vec![worktree_root.clone()]);
664 for locator in locators.iter() {
665 locator.configure(&config);
666 }
667
668 let reporter = pet_reporter::collect::create_reporter();
669 pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
670
671 let mut toolchains = reporter
672 .environments
673 .lock()
674 .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
675
676 let wr = worktree_root;
677 let wr_venv = get_worktree_venv_declaration(&wr);
678 // Sort detected environments by:
679 // environment name matching activation file (<workdir>/.venv)
680 // environment project dir matching worktree_root
681 // general env priority
682 // environment path matching the CONDA_PREFIX env var
683 // executable path
684 toolchains.sort_by(|lhs, rhs| {
685 // Compare venv names against worktree .venv file
686 let venv_ordering =
687 wr_venv
688 .as_ref()
689 .map_or(Ordering::Equal, |venv| match (&lhs.name, &rhs.name) {
690 (Some(l), Some(r)) => (r == venv).cmp(&(l == venv)),
691 (Some(l), None) if l == venv => Ordering::Less,
692 (None, Some(r)) if r == venv => Ordering::Greater,
693 _ => Ordering::Equal,
694 });
695
696 // Compare project paths against worktree root
697 let proj_ordering = || match (&lhs.project, &rhs.project) {
698 (Some(l), Some(r)) => (r == &wr).cmp(&(l == &wr)),
699 (Some(l), None) if l == &wr => Ordering::Less,
700 (None, Some(r)) if r == &wr => Ordering::Greater,
701 _ => Ordering::Equal,
702 };
703
704 // Compare environment priorities
705 let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind));
706
707 // Compare conda prefixes
708 let conda_ordering = || {
709 if lhs.kind == Some(PythonEnvironmentKind::Conda) {
710 environment
711 .get_env_var("CONDA_PREFIX".to_string())
712 .map(|conda_prefix| {
713 let is_match = |exe: &Option<PathBuf>| {
714 exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix))
715 };
716 match (is_match(&lhs.executable), is_match(&rhs.executable)) {
717 (true, false) => Ordering::Less,
718 (false, true) => Ordering::Greater,
719 _ => Ordering::Equal,
720 }
721 })
722 .unwrap_or(Ordering::Equal)
723 } else {
724 Ordering::Equal
725 }
726 };
727
728 // Compare Python executables
729 let exe_ordering = || lhs.executable.cmp(&rhs.executable);
730
731 venv_ordering
732 .then_with(proj_ordering)
733 .then_with(priority_ordering)
734 .then_with(conda_ordering)
735 .then_with(exe_ordering)
736 });
737
738 let mut toolchains: Vec<_> = toolchains
739 .into_iter()
740 .filter_map(|toolchain| {
741 let mut name = String::from("Python");
742 if let Some(ref version) = toolchain.version {
743 _ = write!(name, " {version}");
744 }
745
746 let name_and_kind = match (&toolchain.name, &toolchain.kind) {
747 (Some(name), Some(kind)) => {
748 Some(format!("({name}; {})", python_env_kind_display(kind)))
749 }
750 (Some(name), None) => Some(format!("({name})")),
751 (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
752 (None, None) => None,
753 };
754
755 if let Some(nk) = name_and_kind {
756 _ = write!(name, " {nk}");
757 }
758
759 Some(Toolchain {
760 name: name.into(),
761 path: toolchain.executable.as_ref()?.to_str()?.to_owned().into(),
762 language_name: LanguageName::new("Python"),
763 as_json: serde_json::to_value(toolchain).ok()?,
764 })
765 })
766 .collect();
767 toolchains.dedup();
768 ToolchainList {
769 toolchains,
770 default: None,
771 groups: Default::default(),
772 }
773 }
774 fn term(&self) -> SharedString {
775 self.term.clone()
776 }
777}
778
779pub struct EnvironmentApi<'a> {
780 global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
781 project_env: &'a HashMap<String, String>,
782 pet_env: pet_core::os_environment::EnvironmentApi,
783}
784
785impl<'a> EnvironmentApi<'a> {
786 pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
787 let paths = project_env
788 .get("PATH")
789 .map(|p| std::env::split_paths(p).collect())
790 .unwrap_or_default();
791
792 EnvironmentApi {
793 global_search_locations: Arc::new(Mutex::new(paths)),
794 project_env,
795 pet_env: pet_core::os_environment::EnvironmentApi::new(),
796 }
797 }
798
799 fn user_home(&self) -> Option<PathBuf> {
800 self.project_env
801 .get("HOME")
802 .or_else(|| self.project_env.get("USERPROFILE"))
803 .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
804 .or_else(|| self.pet_env.get_user_home())
805 }
806}
807
808impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
809 fn get_user_home(&self) -> Option<PathBuf> {
810 self.user_home()
811 }
812
813 fn get_root(&self) -> Option<PathBuf> {
814 None
815 }
816
817 fn get_env_var(&self, key: String) -> Option<String> {
818 self.project_env
819 .get(&key)
820 .cloned()
821 .or_else(|| self.pet_env.get_env_var(key))
822 }
823
824 fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
825 if self.global_search_locations.lock().is_empty() {
826 let mut paths =
827 std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
828 .collect::<Vec<PathBuf>>();
829
830 log::trace!("Env PATH: {:?}", paths);
831 for p in self.pet_env.get_know_global_search_locations() {
832 if !paths.contains(&p) {
833 paths.push(p);
834 }
835 }
836
837 let mut paths = paths
838 .into_iter()
839 .filter(|p| p.exists())
840 .collect::<Vec<PathBuf>>();
841
842 self.global_search_locations.lock().append(&mut paths);
843 }
844 self.global_search_locations.lock().clone()
845 }
846}
847
848pub(crate) struct PyLspAdapter {
849 python_venv_base: OnceCell<Result<Arc<Path>, String>>,
850}
851impl PyLspAdapter {
852 const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
853 pub(crate) fn new() -> Self {
854 Self {
855 python_venv_base: OnceCell::new(),
856 }
857 }
858 async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
859 let python_path = Self::find_base_python(delegate)
860 .await
861 .ok_or_else(|| anyhow!("Could not find Python installation for PyLSP"))?;
862 let work_dir = delegate
863 .language_server_download_dir(&Self::SERVER_NAME)
864 .await
865 .ok_or_else(|| anyhow!("Could not get working directory for PyLSP"))?;
866 let mut path = PathBuf::from(work_dir.as_ref());
867 path.push("pylsp-venv");
868 if !path.exists() {
869 util::command::new_smol_command(python_path)
870 .arg("-m")
871 .arg("venv")
872 .arg("pylsp-venv")
873 .current_dir(work_dir)
874 .spawn()?
875 .output()
876 .await?;
877 }
878
879 Ok(path.into())
880 }
881 // Find "baseline", user python version from which we'll create our own venv.
882 async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
883 for path in ["python3", "python"] {
884 if let Some(path) = delegate.which(path.as_ref()).await {
885 return Some(path);
886 }
887 }
888 None
889 }
890
891 async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
892 self.python_venv_base
893 .get_or_init(move || async move {
894 Self::ensure_venv(delegate)
895 .await
896 .map_err(|e| format!("{e}"))
897 })
898 .await
899 .clone()
900 }
901}
902
903const BINARY_DIR: &str = if cfg!(target_os = "windows") {
904 "Scripts"
905} else {
906 "bin"
907};
908
909#[async_trait(?Send)]
910impl LspAdapter for PyLspAdapter {
911 fn name(&self) -> LanguageServerName {
912 Self::SERVER_NAME.clone()
913 }
914
915 async fn check_if_user_installed(
916 &self,
917 delegate: &dyn LspAdapterDelegate,
918 toolchains: Arc<dyn LanguageToolchainStore>,
919 cx: &AsyncApp,
920 ) -> Option<LanguageServerBinary> {
921 if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
922 let env = delegate.shell_env().await;
923 Some(LanguageServerBinary {
924 path: pylsp_bin,
925 env: Some(env),
926 arguments: vec![],
927 })
928 } else {
929 let venv = toolchains
930 .active_toolchain(
931 delegate.worktree_id(),
932 Arc::from("".as_ref()),
933 LanguageName::new("Python"),
934 &mut cx.clone(),
935 )
936 .await?;
937 let pylsp_path = Path::new(venv.path.as_ref()).parent()?.join("pylsp");
938 pylsp_path.exists().then(|| LanguageServerBinary {
939 path: venv.path.to_string().into(),
940 arguments: vec![pylsp_path.into()],
941 env: None,
942 })
943 }
944 }
945
946 async fn fetch_latest_server_version(
947 &self,
948 _: &dyn LspAdapterDelegate,
949 ) -> Result<Box<dyn 'static + Any + Send>> {
950 Ok(Box::new(()) as Box<_>)
951 }
952
953 async fn fetch_server_binary(
954 &self,
955 _: Box<dyn 'static + Send + Any>,
956 _: PathBuf,
957 delegate: &dyn LspAdapterDelegate,
958 ) -> Result<LanguageServerBinary> {
959 let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
960 let pip_path = venv.join(BINARY_DIR).join("pip3");
961 ensure!(
962 util::command::new_smol_command(pip_path.as_path())
963 .arg("install")
964 .arg("python-lsp-server")
965 .arg("-U")
966 .output()
967 .await?
968 .status
969 .success(),
970 "python-lsp-server installation failed"
971 );
972 ensure!(
973 util::command::new_smol_command(pip_path.as_path())
974 .arg("install")
975 .arg("python-lsp-server[all]")
976 .arg("-U")
977 .output()
978 .await?
979 .status
980 .success(),
981 "python-lsp-server[all] installation failed"
982 );
983 ensure!(
984 util::command::new_smol_command(pip_path)
985 .arg("install")
986 .arg("pylsp-mypy")
987 .arg("-U")
988 .output()
989 .await?
990 .status
991 .success(),
992 "pylsp-mypy installation failed"
993 );
994 let pylsp = venv.join(BINARY_DIR).join("pylsp");
995 Ok(LanguageServerBinary {
996 path: pylsp,
997 env: None,
998 arguments: vec![],
999 })
1000 }
1001
1002 async fn cached_server_binary(
1003 &self,
1004 _: PathBuf,
1005 delegate: &dyn LspAdapterDelegate,
1006 ) -> Option<LanguageServerBinary> {
1007 let venv = self.base_venv(delegate).await.ok()?;
1008 let pylsp = venv.join(BINARY_DIR).join("pylsp");
1009 Some(LanguageServerBinary {
1010 path: pylsp,
1011 env: None,
1012 arguments: vec![],
1013 })
1014 }
1015
1016 async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
1017
1018 async fn label_for_completion(
1019 &self,
1020 item: &lsp::CompletionItem,
1021 language: &Arc<language::Language>,
1022 ) -> Option<language::CodeLabel> {
1023 let label = &item.label;
1024 let grammar = language.grammar()?;
1025 let highlight_id = match item.kind? {
1026 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1027 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1028 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1029 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1030 _ => return None,
1031 };
1032 Some(language::CodeLabel {
1033 text: label.clone(),
1034 runs: vec![(0..label.len(), highlight_id)],
1035 filter_range: 0..label.len(),
1036 })
1037 }
1038
1039 async fn label_for_symbol(
1040 &self,
1041 name: &str,
1042 kind: lsp::SymbolKind,
1043 language: &Arc<language::Language>,
1044 ) -> Option<language::CodeLabel> {
1045 let (text, filter_range, display_range) = match kind {
1046 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1047 let text = format!("def {}():\n", name);
1048 let filter_range = 4..4 + name.len();
1049 let display_range = 0..filter_range.end;
1050 (text, filter_range, display_range)
1051 }
1052 lsp::SymbolKind::CLASS => {
1053 let text = format!("class {}:", name);
1054 let filter_range = 6..6 + name.len();
1055 let display_range = 0..filter_range.end;
1056 (text, filter_range, display_range)
1057 }
1058 lsp::SymbolKind::CONSTANT => {
1059 let text = format!("{} = 0", name);
1060 let filter_range = 0..name.len();
1061 let display_range = 0..filter_range.end;
1062 (text, filter_range, display_range)
1063 }
1064 _ => return None,
1065 };
1066
1067 Some(language::CodeLabel {
1068 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
1069 text: text[display_range].to_string(),
1070 filter_range,
1071 })
1072 }
1073
1074 async fn workspace_configuration(
1075 self: Arc<Self>,
1076 _: &dyn Fs,
1077 adapter: &Arc<dyn LspAdapterDelegate>,
1078 toolchains: Arc<dyn LanguageToolchainStore>,
1079 cx: &mut AsyncApp,
1080 ) -> Result<Value> {
1081 let toolchain = toolchains
1082 .active_toolchain(
1083 adapter.worktree_id(),
1084 Arc::from("".as_ref()),
1085 LanguageName::new("Python"),
1086 cx,
1087 )
1088 .await;
1089 cx.update(move |cx| {
1090 let mut user_settings =
1091 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1092 .and_then(|s| s.settings.clone())
1093 .unwrap_or_else(|| {
1094 json!({
1095 "plugins": {
1096 "pycodestyle": {"enabled": false},
1097 "rope_autoimport": {"enabled": true, "memory": true},
1098 "pylsp_mypy": {"enabled": false}
1099 },
1100 "rope": {
1101 "ropeFolder": null
1102 },
1103 })
1104 });
1105
1106 // If user did not explicitly modify their python venv, use one from picker.
1107 if let Some(toolchain) = toolchain {
1108 if user_settings.is_null() {
1109 user_settings = Value::Object(serde_json::Map::default());
1110 }
1111 let object = user_settings.as_object_mut().unwrap();
1112 if let Some(python) = object
1113 .entry("plugins")
1114 .or_insert(Value::Object(serde_json::Map::default()))
1115 .as_object_mut()
1116 {
1117 if let Some(jedi) = python
1118 .entry("jedi")
1119 .or_insert(Value::Object(serde_json::Map::default()))
1120 .as_object_mut()
1121 {
1122 jedi.entry("environment".to_string())
1123 .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1124 }
1125 if let Some(pylint) = python
1126 .entry("pylsp_mypy")
1127 .or_insert(Value::Object(serde_json::Map::default()))
1128 .as_object_mut()
1129 {
1130 pylint.entry("overrides".to_string()).or_insert_with(|| {
1131 Value::Array(vec![
1132 Value::String("--python-executable".into()),
1133 Value::String(toolchain.path.into()),
1134 Value::String("--cache-dir=/dev/null".into()),
1135 Value::Bool(true),
1136 ])
1137 });
1138 }
1139 }
1140 }
1141 user_settings = Value::Object(serde_json::Map::from_iter([(
1142 "pylsp".to_string(),
1143 user_settings,
1144 )]));
1145
1146 user_settings
1147 })
1148 }
1149}
1150
1151#[cfg(test)]
1152mod tests {
1153 use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
1154 use language::{AutoindentMode, Buffer, language_settings::AllLanguageSettings};
1155 use settings::SettingsStore;
1156 use std::num::NonZeroU32;
1157
1158 #[gpui::test]
1159 async fn test_python_autoindent(cx: &mut TestAppContext) {
1160 cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1161 let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1162 cx.update(|cx| {
1163 let test_settings = SettingsStore::test(cx);
1164 cx.set_global(test_settings);
1165 language::init(cx);
1166 cx.update_global::<SettingsStore, _>(|store, cx| {
1167 store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1168 s.defaults.tab_size = NonZeroU32::new(2);
1169 });
1170 });
1171 });
1172
1173 cx.new(|cx| {
1174 let mut buffer = Buffer::local("", cx).with_language(language, cx);
1175 let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
1176 let ix = buffer.len();
1177 buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1178 };
1179
1180 // indent after "def():"
1181 append(&mut buffer, "def a():\n", cx);
1182 assert_eq!(buffer.text(), "def a():\n ");
1183
1184 // preserve indent after blank line
1185 append(&mut buffer, "\n ", cx);
1186 assert_eq!(buffer.text(), "def a():\n \n ");
1187
1188 // indent after "if"
1189 append(&mut buffer, "if a:\n ", cx);
1190 assert_eq!(buffer.text(), "def a():\n \n if a:\n ");
1191
1192 // preserve indent after statement
1193 append(&mut buffer, "b()\n", cx);
1194 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n ");
1195
1196 // preserve indent after statement
1197 append(&mut buffer, "else", cx);
1198 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else");
1199
1200 // dedent "else""
1201 append(&mut buffer, ":", cx);
1202 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else:");
1203
1204 // indent lines after else
1205 append(&mut buffer, "\n", cx);
1206 assert_eq!(
1207 buffer.text(),
1208 "def a():\n \n if a:\n b()\n else:\n "
1209 );
1210
1211 // indent after an open paren. the closing paren is not indented
1212 // because there is another token before it on the same line.
1213 append(&mut buffer, "foo(\n1)", cx);
1214 assert_eq!(
1215 buffer.text(),
1216 "def a():\n \n if a:\n b()\n else:\n foo(\n 1)"
1217 );
1218
1219 // dedent the closing paren if it is shifted to the beginning of the line
1220 let argument_ix = buffer.text().find('1').unwrap();
1221 buffer.edit(
1222 [(argument_ix..argument_ix + 1, "")],
1223 Some(AutoindentMode::EachLine),
1224 cx,
1225 );
1226 assert_eq!(
1227 buffer.text(),
1228 "def a():\n \n if a:\n b()\n else:\n foo(\n )"
1229 );
1230
1231 // preserve indent after the close paren
1232 append(&mut buffer, "\n", cx);
1233 assert_eq!(
1234 buffer.text(),
1235 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n "
1236 );
1237
1238 // manually outdent the last line
1239 let end_whitespace_ix = buffer.len() - 4;
1240 buffer.edit(
1241 [(end_whitespace_ix..buffer.len(), "")],
1242 Some(AutoindentMode::EachLine),
1243 cx,
1244 );
1245 assert_eq!(
1246 buffer.text(),
1247 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n"
1248 );
1249
1250 // preserve the newly reduced indentation on the next newline
1251 append(&mut buffer, "\n", cx);
1252 assert_eq!(
1253 buffer.text(),
1254 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n\n"
1255 );
1256
1257 // reset to a simple if statement
1258 buffer.edit([(0..buffer.len(), "if a:\n b(\n )")], None, cx);
1259
1260 // dedent "else" on the line after a closing paren
1261 append(&mut buffer, "\n else:\n", cx);
1262 assert_eq!(buffer.text(), "if a:\n b(\n )\nelse:\n ");
1263
1264 buffer
1265 });
1266 }
1267}