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