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 std::str::FromStr;
27use std::sync::Mutex;
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 .ok()
675 .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
676
677 let wr = worktree_root;
678 let wr_venv = get_worktree_venv_declaration(&wr);
679 // Sort detected environments by:
680 // environment name matching activation file (<workdir>/.venv)
681 // environment project dir matching worktree_root
682 // general env priority
683 // environment path matching the CONDA_PREFIX env var
684 // executable path
685 toolchains.sort_by(|lhs, rhs| {
686 // Compare venv names against worktree .venv file
687 let venv_ordering =
688 wr_venv
689 .as_ref()
690 .map_or(Ordering::Equal, |venv| match (&lhs.name, &rhs.name) {
691 (Some(l), Some(r)) => (r == venv).cmp(&(l == venv)),
692 (Some(l), None) if l == venv => Ordering::Less,
693 (None, Some(r)) if r == venv => Ordering::Greater,
694 _ => Ordering::Equal,
695 });
696
697 // Compare project paths against worktree root
698 let proj_ordering = || match (&lhs.project, &rhs.project) {
699 (Some(l), Some(r)) => (r == &wr).cmp(&(l == &wr)),
700 (Some(l), None) if l == &wr => Ordering::Less,
701 (None, Some(r)) if r == &wr => Ordering::Greater,
702 _ => Ordering::Equal,
703 };
704
705 // Compare environment priorities
706 let priority_ordering = || env_priority(lhs.kind).cmp(&env_priority(rhs.kind));
707
708 // Compare conda prefixes
709 let conda_ordering = || {
710 if lhs.kind == Some(PythonEnvironmentKind::Conda) {
711 environment
712 .get_env_var("CONDA_PREFIX".to_string())
713 .map(|conda_prefix| {
714 let is_match = |exe: &Option<PathBuf>| {
715 exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix))
716 };
717 match (is_match(&lhs.executable), is_match(&rhs.executable)) {
718 (true, false) => Ordering::Less,
719 (false, true) => Ordering::Greater,
720 _ => Ordering::Equal,
721 }
722 })
723 .unwrap_or(Ordering::Equal)
724 } else {
725 Ordering::Equal
726 }
727 };
728
729 // Compare Python executables
730 let exe_ordering = || lhs.executable.cmp(&rhs.executable);
731
732 venv_ordering
733 .then_with(proj_ordering)
734 .then_with(priority_ordering)
735 .then_with(conda_ordering)
736 .then_with(exe_ordering)
737 });
738
739 let mut toolchains: Vec<_> = toolchains
740 .into_iter()
741 .filter_map(|toolchain| {
742 let mut name = String::from("Python");
743 if let Some(ref version) = toolchain.version {
744 _ = write!(name, " {version}");
745 }
746
747 let name_and_kind = match (&toolchain.name, &toolchain.kind) {
748 (Some(name), Some(kind)) => {
749 Some(format!("({name}; {})", python_env_kind_display(kind)))
750 }
751 (Some(name), None) => Some(format!("({name})")),
752 (None, Some(kind)) => Some(format!("({})", python_env_kind_display(kind))),
753 (None, None) => None,
754 };
755
756 if let Some(nk) = name_and_kind {
757 _ = write!(name, " {nk}");
758 }
759
760 Some(Toolchain {
761 name: name.into(),
762 path: toolchain.executable.as_ref()?.to_str()?.to_owned().into(),
763 language_name: LanguageName::new("Python"),
764 as_json: serde_json::to_value(toolchain).ok()?,
765 })
766 })
767 .collect();
768 toolchains.dedup();
769 ToolchainList {
770 toolchains,
771 default: None,
772 groups: Default::default(),
773 }
774 }
775 fn term(&self) -> SharedString {
776 self.term.clone()
777 }
778}
779
780pub struct EnvironmentApi<'a> {
781 global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
782 project_env: &'a HashMap<String, String>,
783 pet_env: pet_core::os_environment::EnvironmentApi,
784}
785
786impl<'a> EnvironmentApi<'a> {
787 pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
788 let paths = project_env
789 .get("PATH")
790 .map(|p| std::env::split_paths(p).collect())
791 .unwrap_or_default();
792
793 EnvironmentApi {
794 global_search_locations: Arc::new(Mutex::new(paths)),
795 project_env,
796 pet_env: pet_core::os_environment::EnvironmentApi::new(),
797 }
798 }
799
800 fn user_home(&self) -> Option<PathBuf> {
801 self.project_env
802 .get("HOME")
803 .or_else(|| self.project_env.get("USERPROFILE"))
804 .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
805 .or_else(|| self.pet_env.get_user_home())
806 }
807}
808
809impl pet_core::os_environment::Environment for EnvironmentApi<'_> {
810 fn get_user_home(&self) -> Option<PathBuf> {
811 self.user_home()
812 }
813
814 fn get_root(&self) -> Option<PathBuf> {
815 None
816 }
817
818 fn get_env_var(&self, key: String) -> Option<String> {
819 self.project_env
820 .get(&key)
821 .cloned()
822 .or_else(|| self.pet_env.get_env_var(key))
823 }
824
825 fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
826 if self.global_search_locations.lock().unwrap().is_empty() {
827 let mut paths =
828 std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
829 .collect::<Vec<PathBuf>>();
830
831 log::trace!("Env PATH: {:?}", paths);
832 for p in self.pet_env.get_know_global_search_locations() {
833 if !paths.contains(&p) {
834 paths.push(p);
835 }
836 }
837
838 let mut paths = paths
839 .into_iter()
840 .filter(|p| p.exists())
841 .collect::<Vec<PathBuf>>();
842
843 self.global_search_locations
844 .lock()
845 .unwrap()
846 .append(&mut paths);
847 }
848 self.global_search_locations.lock().unwrap().clone()
849 }
850}
851
852pub(crate) struct PyLspAdapter {
853 python_venv_base: OnceCell<Result<Arc<Path>, String>>,
854}
855impl PyLspAdapter {
856 const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
857 pub(crate) fn new() -> Self {
858 Self {
859 python_venv_base: OnceCell::new(),
860 }
861 }
862 async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
863 let python_path = Self::find_base_python(delegate)
864 .await
865 .ok_or_else(|| anyhow!("Could not find Python installation for PyLSP"))?;
866 let work_dir = delegate
867 .language_server_download_dir(&Self::SERVER_NAME)
868 .await
869 .ok_or_else(|| anyhow!("Could not get working directory for PyLSP"))?;
870 let mut path = PathBuf::from(work_dir.as_ref());
871 path.push("pylsp-venv");
872 if !path.exists() {
873 util::command::new_smol_command(python_path)
874 .arg("-m")
875 .arg("venv")
876 .arg("pylsp-venv")
877 .current_dir(work_dir)
878 .spawn()?
879 .output()
880 .await?;
881 }
882
883 Ok(path.into())
884 }
885 // Find "baseline", user python version from which we'll create our own venv.
886 async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
887 for path in ["python3", "python"] {
888 if let Some(path) = delegate.which(path.as_ref()).await {
889 return Some(path);
890 }
891 }
892 None
893 }
894
895 async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
896 self.python_venv_base
897 .get_or_init(move || async move {
898 Self::ensure_venv(delegate)
899 .await
900 .map_err(|e| format!("{e}"))
901 })
902 .await
903 .clone()
904 }
905}
906
907const BINARY_DIR: &str = if cfg!(target_os = "windows") {
908 "Scripts"
909} else {
910 "bin"
911};
912
913#[async_trait(?Send)]
914impl LspAdapter for PyLspAdapter {
915 fn name(&self) -> LanguageServerName {
916 Self::SERVER_NAME.clone()
917 }
918
919 async fn check_if_user_installed(
920 &self,
921 delegate: &dyn LspAdapterDelegate,
922 toolchains: Arc<dyn LanguageToolchainStore>,
923 cx: &AsyncApp,
924 ) -> Option<LanguageServerBinary> {
925 if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
926 let env = delegate.shell_env().await;
927 Some(LanguageServerBinary {
928 path: pylsp_bin,
929 env: Some(env),
930 arguments: vec![],
931 })
932 } else {
933 let venv = toolchains
934 .active_toolchain(
935 delegate.worktree_id(),
936 Arc::from("".as_ref()),
937 LanguageName::new("Python"),
938 &mut cx.clone(),
939 )
940 .await?;
941 let pylsp_path = Path::new(venv.path.as_ref()).parent()?.join("pylsp");
942 pylsp_path.exists().then(|| LanguageServerBinary {
943 path: venv.path.to_string().into(),
944 arguments: vec![pylsp_path.into()],
945 env: None,
946 })
947 }
948 }
949
950 async fn fetch_latest_server_version(
951 &self,
952 _: &dyn LspAdapterDelegate,
953 ) -> Result<Box<dyn 'static + Any + Send>> {
954 Ok(Box::new(()) as Box<_>)
955 }
956
957 async fn fetch_server_binary(
958 &self,
959 _: Box<dyn 'static + Send + Any>,
960 _: PathBuf,
961 delegate: &dyn LspAdapterDelegate,
962 ) -> Result<LanguageServerBinary> {
963 let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
964 let pip_path = venv.join(BINARY_DIR).join("pip3");
965 ensure!(
966 util::command::new_smol_command(pip_path.as_path())
967 .arg("install")
968 .arg("python-lsp-server")
969 .arg("-U")
970 .output()
971 .await?
972 .status
973 .success(),
974 "python-lsp-server installation failed"
975 );
976 ensure!(
977 util::command::new_smol_command(pip_path.as_path())
978 .arg("install")
979 .arg("python-lsp-server[all]")
980 .arg("-U")
981 .output()
982 .await?
983 .status
984 .success(),
985 "python-lsp-server[all] installation failed"
986 );
987 ensure!(
988 util::command::new_smol_command(pip_path)
989 .arg("install")
990 .arg("pylsp-mypy")
991 .arg("-U")
992 .output()
993 .await?
994 .status
995 .success(),
996 "pylsp-mypy installation failed"
997 );
998 let pylsp = venv.join(BINARY_DIR).join("pylsp");
999 Ok(LanguageServerBinary {
1000 path: pylsp,
1001 env: None,
1002 arguments: vec![],
1003 })
1004 }
1005
1006 async fn cached_server_binary(
1007 &self,
1008 _: PathBuf,
1009 delegate: &dyn LspAdapterDelegate,
1010 ) -> Option<LanguageServerBinary> {
1011 let venv = self.base_venv(delegate).await.ok()?;
1012 let pylsp = venv.join(BINARY_DIR).join("pylsp");
1013 Some(LanguageServerBinary {
1014 path: pylsp,
1015 env: None,
1016 arguments: vec![],
1017 })
1018 }
1019
1020 async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
1021
1022 async fn label_for_completion(
1023 &self,
1024 item: &lsp::CompletionItem,
1025 language: &Arc<language::Language>,
1026 ) -> Option<language::CodeLabel> {
1027 let label = &item.label;
1028 let grammar = language.grammar()?;
1029 let highlight_id = match item.kind? {
1030 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
1031 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
1032 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
1033 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
1034 _ => return None,
1035 };
1036 Some(language::CodeLabel {
1037 text: label.clone(),
1038 runs: vec![(0..label.len(), highlight_id)],
1039 filter_range: 0..label.len(),
1040 })
1041 }
1042
1043 async fn label_for_symbol(
1044 &self,
1045 name: &str,
1046 kind: lsp::SymbolKind,
1047 language: &Arc<language::Language>,
1048 ) -> Option<language::CodeLabel> {
1049 let (text, filter_range, display_range) = match kind {
1050 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
1051 let text = format!("def {}():\n", name);
1052 let filter_range = 4..4 + name.len();
1053 let display_range = 0..filter_range.end;
1054 (text, filter_range, display_range)
1055 }
1056 lsp::SymbolKind::CLASS => {
1057 let text = format!("class {}:", name);
1058 let filter_range = 6..6 + name.len();
1059 let display_range = 0..filter_range.end;
1060 (text, filter_range, display_range)
1061 }
1062 lsp::SymbolKind::CONSTANT => {
1063 let text = format!("{} = 0", name);
1064 let filter_range = 0..name.len();
1065 let display_range = 0..filter_range.end;
1066 (text, filter_range, display_range)
1067 }
1068 _ => return None,
1069 };
1070
1071 Some(language::CodeLabel {
1072 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
1073 text: text[display_range].to_string(),
1074 filter_range,
1075 })
1076 }
1077
1078 async fn workspace_configuration(
1079 self: Arc<Self>,
1080 _: &dyn Fs,
1081 adapter: &Arc<dyn LspAdapterDelegate>,
1082 toolchains: Arc<dyn LanguageToolchainStore>,
1083 cx: &mut AsyncApp,
1084 ) -> Result<Value> {
1085 let toolchain = toolchains
1086 .active_toolchain(
1087 adapter.worktree_id(),
1088 Arc::from("".as_ref()),
1089 LanguageName::new("Python"),
1090 cx,
1091 )
1092 .await;
1093 cx.update(move |cx| {
1094 let mut user_settings =
1095 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
1096 .and_then(|s| s.settings.clone())
1097 .unwrap_or_else(|| {
1098 json!({
1099 "plugins": {
1100 "pycodestyle": {"enabled": false},
1101 "rope_autoimport": {"enabled": true, "memory": true},
1102 "pylsp_mypy": {"enabled": false}
1103 },
1104 "rope": {
1105 "ropeFolder": null
1106 },
1107 })
1108 });
1109
1110 // If user did not explicitly modify their python venv, use one from picker.
1111 if let Some(toolchain) = toolchain {
1112 if user_settings.is_null() {
1113 user_settings = Value::Object(serde_json::Map::default());
1114 }
1115 let object = user_settings.as_object_mut().unwrap();
1116 if let Some(python) = object
1117 .entry("plugins")
1118 .or_insert(Value::Object(serde_json::Map::default()))
1119 .as_object_mut()
1120 {
1121 if let Some(jedi) = python
1122 .entry("jedi")
1123 .or_insert(Value::Object(serde_json::Map::default()))
1124 .as_object_mut()
1125 {
1126 jedi.entry("environment".to_string())
1127 .or_insert_with(|| Value::String(toolchain.path.clone().into()));
1128 }
1129 if let Some(pylint) = python
1130 .entry("pylsp_mypy")
1131 .or_insert(Value::Object(serde_json::Map::default()))
1132 .as_object_mut()
1133 {
1134 pylint.entry("overrides".to_string()).or_insert_with(|| {
1135 Value::Array(vec![
1136 Value::String("--python-executable".into()),
1137 Value::String(toolchain.path.into()),
1138 Value::String("--cache-dir=/dev/null".into()),
1139 Value::Bool(true),
1140 ])
1141 });
1142 }
1143 }
1144 }
1145 user_settings = Value::Object(serde_json::Map::from_iter([(
1146 "pylsp".to_string(),
1147 user_settings,
1148 )]));
1149
1150 user_settings
1151 })
1152 }
1153}
1154
1155#[cfg(test)]
1156mod tests {
1157 use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
1158 use language::{AutoindentMode, Buffer, language_settings::AllLanguageSettings};
1159 use settings::SettingsStore;
1160 use std::num::NonZeroU32;
1161
1162 #[gpui::test]
1163 async fn test_python_autoindent(cx: &mut TestAppContext) {
1164 cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1165 let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1166 cx.update(|cx| {
1167 let test_settings = SettingsStore::test(cx);
1168 cx.set_global(test_settings);
1169 language::init(cx);
1170 cx.update_global::<SettingsStore, _>(|store, cx| {
1171 store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1172 s.defaults.tab_size = NonZeroU32::new(2);
1173 });
1174 });
1175 });
1176
1177 cx.new(|cx| {
1178 let mut buffer = Buffer::local("", cx).with_language(language, cx);
1179 let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
1180 let ix = buffer.len();
1181 buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1182 };
1183
1184 // indent after "def():"
1185 append(&mut buffer, "def a():\n", cx);
1186 assert_eq!(buffer.text(), "def a():\n ");
1187
1188 // preserve indent after blank line
1189 append(&mut buffer, "\n ", cx);
1190 assert_eq!(buffer.text(), "def a():\n \n ");
1191
1192 // indent after "if"
1193 append(&mut buffer, "if a:\n ", cx);
1194 assert_eq!(buffer.text(), "def a():\n \n if a:\n ");
1195
1196 // preserve indent after statement
1197 append(&mut buffer, "b()\n", cx);
1198 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n ");
1199
1200 // preserve indent after statement
1201 append(&mut buffer, "else", cx);
1202 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else");
1203
1204 // dedent "else""
1205 append(&mut buffer, ":", cx);
1206 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else:");
1207
1208 // indent lines after else
1209 append(&mut buffer, "\n", cx);
1210 assert_eq!(
1211 buffer.text(),
1212 "def a():\n \n if a:\n b()\n else:\n "
1213 );
1214
1215 // indent after an open paren. the closing paren is not indented
1216 // because there is another token before it on the same line.
1217 append(&mut buffer, "foo(\n1)", cx);
1218 assert_eq!(
1219 buffer.text(),
1220 "def a():\n \n if a:\n b()\n else:\n foo(\n 1)"
1221 );
1222
1223 // dedent the closing paren if it is shifted to the beginning of the line
1224 let argument_ix = buffer.text().find('1').unwrap();
1225 buffer.edit(
1226 [(argument_ix..argument_ix + 1, "")],
1227 Some(AutoindentMode::EachLine),
1228 cx,
1229 );
1230 assert_eq!(
1231 buffer.text(),
1232 "def a():\n \n if a:\n b()\n else:\n foo(\n )"
1233 );
1234
1235 // preserve indent after the close paren
1236 append(&mut buffer, "\n", cx);
1237 assert_eq!(
1238 buffer.text(),
1239 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n "
1240 );
1241
1242 // manually outdent the last line
1243 let end_whitespace_ix = buffer.len() - 4;
1244 buffer.edit(
1245 [(end_whitespace_ix..buffer.len(), "")],
1246 Some(AutoindentMode::EachLine),
1247 cx,
1248 );
1249 assert_eq!(
1250 buffer.text(),
1251 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n"
1252 );
1253
1254 // preserve the newly reduced indentation on the next newline
1255 append(&mut buffer, "\n", cx);
1256 assert_eq!(
1257 buffer.text(),
1258 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n\n"
1259 );
1260
1261 // reset to a simple if statement
1262 buffer.edit([(0..buffer.len(), "if a:\n b(\n )")], None, cx);
1263
1264 // dedent "else" on the line after a closing paren
1265 append(&mut buffer, "\n else:\n", cx);
1266 assert_eq!(buffer.text(), "if a:\n b(\n )\nelse:\n ");
1267
1268 buffer
1269 });
1270 }
1271}