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