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