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 if let Some(pyright_bin) = delegate.which("pyright-langserver".as_ref()).await {
87 let env = delegate.shell_env().await;
88 Some(LanguageServerBinary {
89 path: pyright_bin,
90 env: Some(env),
91 arguments: vec!["--stdio".into()],
92 })
93 } else {
94 let node = delegate.which("node".as_ref()).await?;
95 let (node_modules_path, _) = delegate
96 .npm_package_installed_version(Self::SERVER_NAME.as_ref())
97 .await
98 .log_err()??;
99
100 let path = node_modules_path.join(NODE_MODULE_RELATIVE_SERVER_PATH);
101
102 Some(LanguageServerBinary {
103 path: node,
104 env: None,
105 arguments: server_binary_arguments(&path),
106 })
107 }
108 }
109
110 async fn fetch_latest_server_version(
111 &self,
112 _: &dyn LspAdapterDelegate,
113 ) -> Result<Box<dyn 'static + Any + Send>> {
114 Ok(Box::new(
115 self.node
116 .npm_package_latest_version(Self::SERVER_NAME.as_ref())
117 .await?,
118 ) as Box<_>)
119 }
120
121 async fn fetch_server_binary(
122 &self,
123 latest_version: Box<dyn 'static + Send + Any>,
124 container_dir: PathBuf,
125 _: &dyn LspAdapterDelegate,
126 ) -> Result<LanguageServerBinary> {
127 let latest_version = latest_version.downcast::<String>().unwrap();
128 let server_path = container_dir.join(SERVER_PATH);
129
130 self.node
131 .npm_install_packages(
132 &container_dir,
133 &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
134 )
135 .await?;
136
137 Ok(LanguageServerBinary {
138 path: self.node.binary_path().await?,
139 env: None,
140 arguments: server_binary_arguments(&server_path),
141 })
142 }
143
144 async fn check_if_version_installed(
145 &self,
146 version: &(dyn 'static + Send + Any),
147 container_dir: &PathBuf,
148 _: &dyn LspAdapterDelegate,
149 ) -> Option<LanguageServerBinary> {
150 let version = version.downcast_ref::<String>().unwrap();
151 let server_path = container_dir.join(SERVER_PATH);
152
153 let should_install_language_server = self
154 .node
155 .should_install_npm_package(
156 Self::SERVER_NAME.as_ref(),
157 &server_path,
158 &container_dir,
159 &version,
160 )
161 .await;
162
163 if should_install_language_server {
164 None
165 } else {
166 Some(LanguageServerBinary {
167 path: self.node.binary_path().await.ok()?,
168 env: None,
169 arguments: server_binary_arguments(&server_path),
170 })
171 }
172 }
173
174 async fn cached_server_binary(
175 &self,
176 container_dir: PathBuf,
177 _: &dyn LspAdapterDelegate,
178 ) -> Option<LanguageServerBinary> {
179 get_cached_server_binary(container_dir, &self.node).await
180 }
181
182 async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
183 // Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
184 // Where `XX` is the sorting category, `YYYY` is based on most recent usage,
185 // and `name` is the symbol name itself.
186 //
187 // Because the symbol name is included, there generally are not ties when
188 // sorting by the `sortText`, so the symbol's fuzzy match score is not taken
189 // into account. Here, we remove the symbol name from the sortText in order
190 // to allow our own fuzzy score to be used to break ties.
191 //
192 // see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
193 for item in items {
194 let Some(sort_text) = &mut item.sort_text else {
195 continue;
196 };
197 let mut parts = sort_text.split('.');
198 let Some(first) = parts.next() else { continue };
199 let Some(second) = parts.next() else { continue };
200 let Some(_) = parts.next() else { continue };
201 sort_text.replace_range(first.len() + second.len() + 1.., "");
202 }
203 }
204
205 async fn label_for_completion(
206 &self,
207 item: &lsp::CompletionItem,
208 language: &Arc<language::Language>,
209 ) -> Option<language::CodeLabel> {
210 let label = &item.label;
211 let grammar = language.grammar()?;
212 let highlight_id = match item.kind? {
213 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
214 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
215 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
216 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
217 _ => return None,
218 };
219 Some(language::CodeLabel {
220 text: label.clone(),
221 runs: vec![(0..label.len(), highlight_id)],
222 filter_range: 0..label.len(),
223 })
224 }
225
226 async fn label_for_symbol(
227 &self,
228 name: &str,
229 kind: lsp::SymbolKind,
230 language: &Arc<language::Language>,
231 ) -> Option<language::CodeLabel> {
232 let (text, filter_range, display_range) = match kind {
233 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
234 let text = format!("def {}():\n", name);
235 let filter_range = 4..4 + name.len();
236 let display_range = 0..filter_range.end;
237 (text, filter_range, display_range)
238 }
239 lsp::SymbolKind::CLASS => {
240 let text = format!("class {}:", name);
241 let filter_range = 6..6 + name.len();
242 let display_range = 0..filter_range.end;
243 (text, filter_range, display_range)
244 }
245 lsp::SymbolKind::CONSTANT => {
246 let text = format!("{} = 0", name);
247 let filter_range = 0..name.len();
248 let display_range = 0..filter_range.end;
249 (text, filter_range, display_range)
250 }
251 _ => return None,
252 };
253
254 Some(language::CodeLabel {
255 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
256 text: text[display_range].to_string(),
257 filter_range,
258 })
259 }
260
261 async fn workspace_configuration(
262 self: Arc<Self>,
263 _: &dyn Fs,
264 adapter: &Arc<dyn LspAdapterDelegate>,
265 toolchains: Arc<dyn LanguageToolchainStore>,
266 cx: &mut AsyncApp,
267 ) -> Result<Value> {
268 let toolchain = toolchains
269 .active_toolchain(adapter.worktree_id(), LanguageName::new("Python"), cx)
270 .await;
271 cx.update(move |cx| {
272 let mut user_settings =
273 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
274 .and_then(|s| s.settings.clone())
275 .unwrap_or_default();
276
277 // If python.pythonPath is not set in user config, do so using our toolchain picker.
278 if let Some(toolchain) = toolchain {
279 if user_settings.is_null() {
280 user_settings = Value::Object(serde_json::Map::default());
281 }
282 let object = user_settings.as_object_mut().unwrap();
283 if let Some(python) = object
284 .entry("python")
285 .or_insert(Value::Object(serde_json::Map::default()))
286 .as_object_mut()
287 {
288 python
289 .entry("pythonPath")
290 .or_insert(Value::String(toolchain.path.into()));
291 }
292 }
293 user_settings
294 })
295 }
296}
297
298async fn get_cached_server_binary(
299 container_dir: PathBuf,
300 node: &NodeRuntime,
301) -> Option<LanguageServerBinary> {
302 let server_path = container_dir.join(SERVER_PATH);
303 if server_path.exists() {
304 Some(LanguageServerBinary {
305 path: node.binary_path().await.log_err()?,
306 env: None,
307 arguments: server_binary_arguments(&server_path),
308 })
309 } else {
310 log::error!("missing executable in directory {:?}", server_path);
311 None
312 }
313}
314
315pub(crate) struct PythonContextProvider;
316
317const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
318 VariableName::Custom(Cow::Borrowed("PYTHON_TEST_TARGET"));
319
320const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
321 VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
322impl ContextProvider for PythonContextProvider {
323 fn build_context(
324 &self,
325 variables: &task::TaskVariables,
326 location: &project::Location,
327 _: Option<HashMap<String, String>>,
328 toolchains: Arc<dyn LanguageToolchainStore>,
329 cx: &mut gpui::App,
330 ) -> Task<Result<task::TaskVariables>> {
331 let test_target = match selected_test_runner(location.buffer.read(cx).file(), cx) {
332 TestRunner::UNITTEST => self.build_unittest_target(variables),
333 TestRunner::PYTEST => self.build_pytest_target(variables),
334 };
335
336 let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
337 cx.spawn(async move |cx| {
338 let active_toolchain = if let Some(worktree_id) = worktree_id {
339 toolchains
340 .active_toolchain(worktree_id, "Python".into(), cx)
341 .await
342 .map_or_else(
343 || "python3".to_owned(),
344 |toolchain| format!("\"{}\"", toolchain.path),
345 )
346 } else {
347 String::from("python3")
348 };
349 let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
350 Ok(task::TaskVariables::from_iter(
351 test_target.into_iter().chain([toolchain]),
352 ))
353 })
354 }
355
356 fn associated_tasks(
357 &self,
358 file: Option<Arc<dyn language::File>>,
359 cx: &App,
360 ) -> Option<TaskTemplates> {
361 let test_runner = selected_test_runner(file.as_ref(), cx);
362
363 let mut tasks = vec![
364 // Execute a selection
365 TaskTemplate {
366 label: "execute selection".to_owned(),
367 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
368 args: vec![
369 "-c".to_owned(),
370 VariableName::SelectedText.template_value_with_whitespace(),
371 ],
372 ..TaskTemplate::default()
373 },
374 // Execute an entire file
375 TaskTemplate {
376 label: format!("run '{}'", VariableName::File.template_value()),
377 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
378 args: vec![VariableName::File.template_value_with_whitespace()],
379 ..TaskTemplate::default()
380 },
381 ];
382
383 tasks.extend(match test_runner {
384 TestRunner::UNITTEST => {
385 [
386 // Run tests for an entire file
387 TaskTemplate {
388 label: format!("unittest '{}'", VariableName::File.template_value()),
389 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
390 args: vec![
391 "-m".to_owned(),
392 "unittest".to_owned(),
393 VariableName::File.template_value_with_whitespace(),
394 ],
395 ..TaskTemplate::default()
396 },
397 // Run test(s) for a specific target within a file
398 TaskTemplate {
399 label: "unittest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
400 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
401 args: vec![
402 "-m".to_owned(),
403 "unittest".to_owned(),
404 PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
405 ],
406 tags: vec![
407 "python-unittest-class".to_owned(),
408 "python-unittest-method".to_owned(),
409 ],
410 ..TaskTemplate::default()
411 },
412 ]
413 }
414 TestRunner::PYTEST => {
415 [
416 // Run tests for an entire file
417 TaskTemplate {
418 label: format!("pytest '{}'", VariableName::File.template_value()),
419 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
420 args: vec![
421 "-m".to_owned(),
422 "pytest".to_owned(),
423 VariableName::File.template_value_with_whitespace(),
424 ],
425 ..TaskTemplate::default()
426 },
427 // Run test(s) for a specific target within a file
428 TaskTemplate {
429 label: "pytest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
430 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
431 args: vec![
432 "-m".to_owned(),
433 "pytest".to_owned(),
434 PYTHON_TEST_TARGET_TASK_VARIABLE.template_value_with_whitespace(),
435 ],
436 tags: vec![
437 "python-pytest-class".to_owned(),
438 "python-pytest-method".to_owned(),
439 ],
440 ..TaskTemplate::default()
441 },
442 ]
443 }
444 });
445
446 Some(TaskTemplates(tasks))
447 }
448}
449
450fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &App) -> TestRunner {
451 const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
452 language_settings(Some(LanguageName::new("Python")), location, cx)
453 .tasks
454 .variables
455 .get(TEST_RUNNER_VARIABLE)
456 .and_then(|val| TestRunner::from_str(val).ok())
457 .unwrap_or(TestRunner::PYTEST)
458}
459
460impl PythonContextProvider {
461 fn build_unittest_target(
462 &self,
463 variables: &task::TaskVariables,
464 ) -> Option<(VariableName, String)> {
465 let python_module_name =
466 python_module_name_from_relative_path(variables.get(&VariableName::RelativeFile)?);
467
468 let unittest_class_name =
469 variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
470
471 let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
472 "_unittest_method_name",
473 )));
474
475 let unittest_target_str = match (unittest_class_name, unittest_method_name) {
476 (Some(class_name), Some(method_name)) => {
477 format!("{python_module_name}.{class_name}.{method_name}")
478 }
479 (Some(class_name), None) => format!("{python_module_name}.{class_name}"),
480 (None, None) => python_module_name,
481 // should never happen, a TestCase class is the unit of testing
482 (None, Some(_)) => return None,
483 };
484
485 Some((
486 PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
487 unittest_target_str,
488 ))
489 }
490
491 fn build_pytest_target(
492 &self,
493 variables: &task::TaskVariables,
494 ) -> Option<(VariableName, String)> {
495 let file_path = variables.get(&VariableName::RelativeFile)?;
496
497 let pytest_class_name =
498 variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
499
500 let pytest_method_name =
501 variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
502
503 let pytest_target_str = match (pytest_class_name, pytest_method_name) {
504 (Some(class_name), Some(method_name)) => {
505 format!("{file_path}::{class_name}::{method_name}")
506 }
507 (Some(class_name), None) => {
508 format!("{file_path}::{class_name}")
509 }
510 (None, Some(method_name)) => {
511 format!("{file_path}::{method_name}")
512 }
513 (None, None) => file_path.to_string(),
514 };
515
516 Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
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 pet_core::os_environment::Environment for EnvironmentApi<'_> {
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 if let Some(pylsp_bin) = delegate.which(Self::SERVER_NAME.as_ref()).await {
795 let env = delegate.shell_env().await;
796 Some(LanguageServerBinary {
797 path: pylsp_bin,
798 env: Some(env),
799 arguments: vec![],
800 })
801 } else {
802 let venv = toolchains
803 .active_toolchain(
804 delegate.worktree_id(),
805 LanguageName::new("Python"),
806 &mut cx.clone(),
807 )
808 .await?;
809 let pylsp_path = Path::new(venv.path.as_ref()).parent()?.join("pylsp");
810 pylsp_path.exists().then(|| LanguageServerBinary {
811 path: venv.path.to_string().into(),
812 arguments: vec![pylsp_path.into()],
813 env: None,
814 })
815 }
816 }
817
818 async fn fetch_latest_server_version(
819 &self,
820 _: &dyn LspAdapterDelegate,
821 ) -> Result<Box<dyn 'static + Any + Send>> {
822 Ok(Box::new(()) as Box<_>)
823 }
824
825 async fn fetch_server_binary(
826 &self,
827 _: Box<dyn 'static + Send + Any>,
828 _: PathBuf,
829 delegate: &dyn LspAdapterDelegate,
830 ) -> Result<LanguageServerBinary> {
831 let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
832 let pip_path = venv.join(BINARY_DIR).join("pip3");
833 ensure!(
834 util::command::new_smol_command(pip_path.as_path())
835 .arg("install")
836 .arg("python-lsp-server")
837 .output()
838 .await?
839 .status
840 .success(),
841 "python-lsp-server installation failed"
842 );
843 ensure!(
844 util::command::new_smol_command(pip_path.as_path())
845 .arg("install")
846 .arg("python-lsp-server[all]")
847 .output()
848 .await?
849 .status
850 .success(),
851 "python-lsp-server[all] installation failed"
852 );
853 ensure!(
854 util::command::new_smol_command(pip_path)
855 .arg("install")
856 .arg("pylsp-mypy")
857 .output()
858 .await?
859 .status
860 .success(),
861 "pylsp-mypy installation failed"
862 );
863 let pylsp = venv.join(BINARY_DIR).join("pylsp");
864 Ok(LanguageServerBinary {
865 path: pylsp,
866 env: None,
867 arguments: vec![],
868 })
869 }
870
871 async fn cached_server_binary(
872 &self,
873 _: PathBuf,
874 delegate: &dyn LspAdapterDelegate,
875 ) -> Option<LanguageServerBinary> {
876 let venv = self.base_venv(delegate).await.ok()?;
877 let pylsp = venv.join(BINARY_DIR).join("pylsp");
878 Some(LanguageServerBinary {
879 path: pylsp,
880 env: None,
881 arguments: vec![],
882 })
883 }
884
885 async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
886
887 async fn label_for_completion(
888 &self,
889 item: &lsp::CompletionItem,
890 language: &Arc<language::Language>,
891 ) -> Option<language::CodeLabel> {
892 let label = &item.label;
893 let grammar = language.grammar()?;
894 let highlight_id = match item.kind? {
895 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
896 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
897 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
898 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
899 _ => return None,
900 };
901 Some(language::CodeLabel {
902 text: label.clone(),
903 runs: vec![(0..label.len(), highlight_id)],
904 filter_range: 0..label.len(),
905 })
906 }
907
908 async fn label_for_symbol(
909 &self,
910 name: &str,
911 kind: lsp::SymbolKind,
912 language: &Arc<language::Language>,
913 ) -> Option<language::CodeLabel> {
914 let (text, filter_range, display_range) = match kind {
915 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
916 let text = format!("def {}():\n", name);
917 let filter_range = 4..4 + name.len();
918 let display_range = 0..filter_range.end;
919 (text, filter_range, display_range)
920 }
921 lsp::SymbolKind::CLASS => {
922 let text = format!("class {}:", name);
923 let filter_range = 6..6 + name.len();
924 let display_range = 0..filter_range.end;
925 (text, filter_range, display_range)
926 }
927 lsp::SymbolKind::CONSTANT => {
928 let text = format!("{} = 0", name);
929 let filter_range = 0..name.len();
930 let display_range = 0..filter_range.end;
931 (text, filter_range, display_range)
932 }
933 _ => return None,
934 };
935
936 Some(language::CodeLabel {
937 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
938 text: text[display_range].to_string(),
939 filter_range,
940 })
941 }
942
943 async fn workspace_configuration(
944 self: Arc<Self>,
945 _: &dyn Fs,
946 adapter: &Arc<dyn LspAdapterDelegate>,
947 toolchains: Arc<dyn LanguageToolchainStore>,
948 cx: &mut AsyncApp,
949 ) -> Result<Value> {
950 let toolchain = toolchains
951 .active_toolchain(adapter.worktree_id(), LanguageName::new("Python"), cx)
952 .await;
953 cx.update(move |cx| {
954 let mut user_settings =
955 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
956 .and_then(|s| s.settings.clone())
957 .unwrap_or_else(|| {
958 json!({
959 "plugins": {
960 "pycodestyle": {"enabled": false},
961 "rope_autoimport": {"enabled": true, "memory": true},
962 "pylsp_mypy": {"enabled": false}
963 },
964 "rope": {
965 "ropeFolder": null
966 },
967 })
968 });
969
970 // If user did not explicitly modify their python venv, use one from picker.
971 if let Some(toolchain) = toolchain {
972 if user_settings.is_null() {
973 user_settings = Value::Object(serde_json::Map::default());
974 }
975 let object = user_settings.as_object_mut().unwrap();
976 if let Some(python) = object
977 .entry("plugins")
978 .or_insert(Value::Object(serde_json::Map::default()))
979 .as_object_mut()
980 {
981 if let Some(jedi) = python
982 .entry("jedi")
983 .or_insert(Value::Object(serde_json::Map::default()))
984 .as_object_mut()
985 {
986 jedi.entry("environment".to_string())
987 .or_insert_with(|| Value::String(toolchain.path.clone().into()));
988 }
989 if let Some(pylint) = python
990 .entry("pylsp_mypy")
991 .or_insert(Value::Object(serde_json::Map::default()))
992 .as_object_mut()
993 {
994 pylint.entry("overrides".to_string()).or_insert_with(|| {
995 Value::Array(vec![
996 Value::String("--python-executable".into()),
997 Value::String(toolchain.path.into()),
998 Value::String("--cache-dir=/dev/null".into()),
999 Value::Bool(true),
1000 ])
1001 });
1002 }
1003 }
1004 }
1005 user_settings = Value::Object(serde_json::Map::from_iter([(
1006 "pylsp".to_string(),
1007 user_settings,
1008 )]));
1009
1010 user_settings
1011 })
1012 }
1013}
1014
1015#[cfg(test)]
1016mod tests {
1017 use gpui::{AppContext as _, BorrowAppContext, Context, TestAppContext};
1018 use language::{language_settings::AllLanguageSettings, AutoindentMode, Buffer};
1019 use settings::SettingsStore;
1020 use std::num::NonZeroU32;
1021
1022 #[gpui::test]
1023 async fn test_python_autoindent(cx: &mut TestAppContext) {
1024 cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1025 let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1026 cx.update(|cx| {
1027 let test_settings = SettingsStore::test(cx);
1028 cx.set_global(test_settings);
1029 language::init(cx);
1030 cx.update_global::<SettingsStore, _>(|store, cx| {
1031 store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1032 s.defaults.tab_size = NonZeroU32::new(2);
1033 });
1034 });
1035 });
1036
1037 cx.new(|cx| {
1038 let mut buffer = Buffer::local("", cx).with_language(language, cx);
1039 let append = |buffer: &mut Buffer, text: &str, cx: &mut Context<Buffer>| {
1040 let ix = buffer.len();
1041 buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1042 };
1043
1044 // indent after "def():"
1045 append(&mut buffer, "def a():\n", cx);
1046 assert_eq!(buffer.text(), "def a():\n ");
1047
1048 // preserve indent after blank line
1049 append(&mut buffer, "\n ", cx);
1050 assert_eq!(buffer.text(), "def a():\n \n ");
1051
1052 // indent after "if"
1053 append(&mut buffer, "if a:\n ", cx);
1054 assert_eq!(buffer.text(), "def a():\n \n if a:\n ");
1055
1056 // preserve indent after statement
1057 append(&mut buffer, "b()\n", cx);
1058 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n ");
1059
1060 // preserve indent after statement
1061 append(&mut buffer, "else", cx);
1062 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else");
1063
1064 // dedent "else""
1065 append(&mut buffer, ":", cx);
1066 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else:");
1067
1068 // indent lines after else
1069 append(&mut buffer, "\n", cx);
1070 assert_eq!(
1071 buffer.text(),
1072 "def a():\n \n if a:\n b()\n else:\n "
1073 );
1074
1075 // indent after an open paren. the closing paren is not indented
1076 // because there is another token before it on the same line.
1077 append(&mut buffer, "foo(\n1)", cx);
1078 assert_eq!(
1079 buffer.text(),
1080 "def a():\n \n if a:\n b()\n else:\n foo(\n 1)"
1081 );
1082
1083 // dedent the closing paren if it is shifted to the beginning of the line
1084 let argument_ix = buffer.text().find('1').unwrap();
1085 buffer.edit(
1086 [(argument_ix..argument_ix + 1, "")],
1087 Some(AutoindentMode::EachLine),
1088 cx,
1089 );
1090 assert_eq!(
1091 buffer.text(),
1092 "def a():\n \n if a:\n b()\n else:\n foo(\n )"
1093 );
1094
1095 // preserve indent after the close paren
1096 append(&mut buffer, "\n", cx);
1097 assert_eq!(
1098 buffer.text(),
1099 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n "
1100 );
1101
1102 // manually outdent the last line
1103 let end_whitespace_ix = buffer.len() - 4;
1104 buffer.edit(
1105 [(end_whitespace_ix..buffer.len(), "")],
1106 Some(AutoindentMode::EachLine),
1107 cx,
1108 );
1109 assert_eq!(
1110 buffer.text(),
1111 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n"
1112 );
1113
1114 // preserve the newly reduced indentation on the next newline
1115 append(&mut buffer, "\n", cx);
1116 assert_eq!(
1117 buffer.text(),
1118 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n\n"
1119 );
1120
1121 // reset to a simple if statement
1122 buffer.edit([(0..buffer.len(), "if a:\n b(\n )")], None, cx);
1123
1124 // dedent "else" on the line after a closing paren
1125 append(&mut buffer, "\n else:\n", cx);
1126 assert_eq!(buffer.text(), "if a:\n b(\n )\nelse:\n ");
1127
1128 buffer
1129 });
1130 }
1131}