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::Conda,
546 PythonEnvironmentKind::Pyenv,
547 PythonEnvironmentKind::GlobalPaths,
548 PythonEnvironmentKind::Homebrew,
549];
550
551fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
552 if let Some(kind) = kind {
553 ENV_PRIORITY_LIST
554 .iter()
555 .position(|blessed_env| blessed_env == &kind)
556 .unwrap_or(ENV_PRIORITY_LIST.len())
557 } else {
558 // Unknown toolchains are less useful than non-blessed ones.
559 ENV_PRIORITY_LIST.len() + 1
560 }
561}
562
563#[async_trait]
564impl ToolchainLister for PythonToolchainProvider {
565 async fn list(
566 &self,
567 worktree_root: PathBuf,
568 project_env: Option<HashMap<String, String>>,
569 ) -> ToolchainList {
570 let env = project_env.unwrap_or_default();
571 let environment = EnvironmentApi::from_env(&env);
572 let locators = pet::locators::create_locators(
573 Arc::new(pet_conda::Conda::from(&environment)),
574 Arc::new(pet_poetry::Poetry::from(&environment)),
575 &environment,
576 );
577 let mut config = Configuration::default();
578 config.workspace_directories = Some(vec![worktree_root]);
579 for locator in locators.iter() {
580 locator.configure(&config);
581 }
582
583 let reporter = pet_reporter::collect::create_reporter();
584 pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
585
586 let mut toolchains = reporter
587 .environments
588 .lock()
589 .ok()
590 .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
591
592 toolchains.sort_by(|lhs, rhs| {
593 env_priority(lhs.kind)
594 .cmp(&env_priority(rhs.kind))
595 .then_with(|| {
596 if lhs.kind == Some(PythonEnvironmentKind::Conda) {
597 environment
598 .get_env_var("CONDA_PREFIX".to_string())
599 .map(|conda_prefix| {
600 let is_match = |exe: &Option<PathBuf>| {
601 exe.as_ref().map_or(false, |e| e.starts_with(&conda_prefix))
602 };
603 match (is_match(&lhs.executable), is_match(&rhs.executable)) {
604 (true, false) => Ordering::Less,
605 (false, true) => Ordering::Greater,
606 _ => Ordering::Equal,
607 }
608 })
609 .unwrap_or(Ordering::Equal)
610 } else {
611 Ordering::Equal
612 }
613 })
614 .then_with(|| lhs.executable.cmp(&rhs.executable))
615 });
616
617 let mut toolchains: Vec<_> = toolchains
618 .into_iter()
619 .filter_map(|toolchain| {
620 let name = if let Some(version) = &toolchain.version {
621 format!("Python {version} ({:?})", toolchain.kind?)
622 } else {
623 format!("{:?}", toolchain.kind?)
624 }
625 .into();
626 Some(Toolchain {
627 name,
628 path: toolchain.executable.as_ref()?.to_str()?.to_owned().into(),
629 language_name: LanguageName::new("Python"),
630 as_json: serde_json::to_value(toolchain).ok()?,
631 })
632 })
633 .collect();
634 toolchains.dedup();
635 ToolchainList {
636 toolchains,
637 default: None,
638 groups: Default::default(),
639 }
640 }
641 fn term(&self) -> SharedString {
642 self.term.clone()
643 }
644}
645
646pub struct EnvironmentApi<'a> {
647 global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
648 project_env: &'a HashMap<String, String>,
649 pet_env: pet_core::os_environment::EnvironmentApi,
650}
651
652impl<'a> EnvironmentApi<'a> {
653 pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
654 let paths = project_env
655 .get("PATH")
656 .map(|p| std::env::split_paths(p).collect())
657 .unwrap_or_default();
658
659 EnvironmentApi {
660 global_search_locations: Arc::new(Mutex::new(paths)),
661 project_env,
662 pet_env: pet_core::os_environment::EnvironmentApi::new(),
663 }
664 }
665
666 fn user_home(&self) -> Option<PathBuf> {
667 self.project_env
668 .get("HOME")
669 .or_else(|| self.project_env.get("USERPROFILE"))
670 .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
671 .or_else(|| self.pet_env.get_user_home())
672 }
673}
674
675impl<'a> pet_core::os_environment::Environment for EnvironmentApi<'a> {
676 fn get_user_home(&self) -> Option<PathBuf> {
677 self.user_home()
678 }
679
680 fn get_root(&self) -> Option<PathBuf> {
681 None
682 }
683
684 fn get_env_var(&self, key: String) -> Option<String> {
685 self.project_env
686 .get(&key)
687 .cloned()
688 .or_else(|| self.pet_env.get_env_var(key))
689 }
690
691 fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
692 if self.global_search_locations.lock().unwrap().is_empty() {
693 let mut paths =
694 std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
695 .collect::<Vec<PathBuf>>();
696
697 log::trace!("Env PATH: {:?}", paths);
698 for p in self.pet_env.get_know_global_search_locations() {
699 if !paths.contains(&p) {
700 paths.push(p);
701 }
702 }
703
704 let mut paths = paths
705 .into_iter()
706 .filter(|p| p.exists())
707 .collect::<Vec<PathBuf>>();
708
709 self.global_search_locations
710 .lock()
711 .unwrap()
712 .append(&mut paths);
713 }
714 self.global_search_locations.lock().unwrap().clone()
715 }
716}
717
718pub(crate) struct PyLspAdapter {
719 python_venv_base: OnceCell<Result<Arc<Path>, String>>,
720}
721impl PyLspAdapter {
722 const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pylsp");
723 pub(crate) fn new() -> Self {
724 Self {
725 python_venv_base: OnceCell::new(),
726 }
727 }
728 async fn ensure_venv(delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>> {
729 let python_path = Self::find_base_python(delegate)
730 .await
731 .ok_or_else(|| anyhow!("Could not find Python installation for PyLSP"))?;
732 let work_dir = delegate
733 .language_server_download_dir(&Self::SERVER_NAME)
734 .await
735 .ok_or_else(|| anyhow!("Could not get working directory for PyLSP"))?;
736 let mut path = PathBuf::from(work_dir.as_ref());
737 path.push("pylsp-venv");
738 if !path.exists() {
739 util::command::new_smol_command(python_path)
740 .arg("-m")
741 .arg("venv")
742 .arg("pylsp-venv")
743 .current_dir(work_dir)
744 .spawn()?
745 .output()
746 .await?;
747 }
748
749 Ok(path.into())
750 }
751 // Find "baseline", user python version from which we'll create our own venv.
752 async fn find_base_python(delegate: &dyn LspAdapterDelegate) -> Option<PathBuf> {
753 for path in ["python3", "python"] {
754 if let Some(path) = delegate.which(path.as_ref()).await {
755 return Some(path);
756 }
757 }
758 None
759 }
760
761 async fn base_venv(&self, delegate: &dyn LspAdapterDelegate) -> Result<Arc<Path>, String> {
762 self.python_venv_base
763 .get_or_init(move || async move {
764 Self::ensure_venv(delegate)
765 .await
766 .map_err(|e| format!("{e}"))
767 })
768 .await
769 .clone()
770 }
771}
772
773const BINARY_DIR: &str = if cfg!(target_os = "windows") {
774 "Scripts"
775} else {
776 "bin"
777};
778
779#[async_trait(?Send)]
780impl LspAdapter for PyLspAdapter {
781 fn name(&self) -> LanguageServerName {
782 Self::SERVER_NAME.clone()
783 }
784
785 async fn check_if_user_installed(
786 &self,
787 delegate: &dyn LspAdapterDelegate,
788 toolchains: Arc<dyn LanguageToolchainStore>,
789 cx: &AsyncAppContext,
790 ) -> Option<LanguageServerBinary> {
791 let venv = toolchains
792 .active_toolchain(
793 delegate.worktree_id(),
794 LanguageName::new("Python"),
795 &mut cx.clone(),
796 )
797 .await?;
798 let pylsp_path = Path::new(venv.path.as_ref()).parent()?.join("pylsp");
799 pylsp_path.exists().then(|| LanguageServerBinary {
800 path: venv.path.to_string().into(),
801 arguments: vec![pylsp_path.into()],
802 env: None,
803 })
804 }
805
806 async fn fetch_latest_server_version(
807 &self,
808 _: &dyn LspAdapterDelegate,
809 ) -> Result<Box<dyn 'static + Any + Send>> {
810 Ok(Box::new(()) as Box<_>)
811 }
812
813 async fn fetch_server_binary(
814 &self,
815 _: Box<dyn 'static + Send + Any>,
816 _: PathBuf,
817 delegate: &dyn LspAdapterDelegate,
818 ) -> Result<LanguageServerBinary> {
819 let venv = self.base_venv(delegate).await.map_err(|e| anyhow!(e))?;
820 let pip_path = venv.join(BINARY_DIR).join("pip3");
821 ensure!(
822 util::command::new_smol_command(pip_path.as_path())
823 .arg("install")
824 .arg("python-lsp-server")
825 .output()
826 .await?
827 .status
828 .success(),
829 "python-lsp-server installation failed"
830 );
831 ensure!(
832 util::command::new_smol_command(pip_path.as_path())
833 .arg("install")
834 .arg("python-lsp-server[all]")
835 .output()
836 .await?
837 .status
838 .success(),
839 "python-lsp-server[all] installation failed"
840 );
841 ensure!(
842 util::command::new_smol_command(pip_path)
843 .arg("install")
844 .arg("pylsp-mypy")
845 .output()
846 .await?
847 .status
848 .success(),
849 "pylsp-mypy installation failed"
850 );
851 let pylsp = venv.join(BINARY_DIR).join("pylsp");
852 Ok(LanguageServerBinary {
853 path: pylsp,
854 env: None,
855 arguments: vec![],
856 })
857 }
858
859 async fn cached_server_binary(
860 &self,
861 _: PathBuf,
862 delegate: &dyn LspAdapterDelegate,
863 ) -> Option<LanguageServerBinary> {
864 let venv = self.base_venv(delegate).await.ok()?;
865 let pylsp = venv.join(BINARY_DIR).join("pylsp");
866 Some(LanguageServerBinary {
867 path: pylsp,
868 env: None,
869 arguments: vec![],
870 })
871 }
872
873 async fn process_completions(&self, _items: &mut [lsp::CompletionItem]) {}
874
875 async fn label_for_completion(
876 &self,
877 item: &lsp::CompletionItem,
878 language: &Arc<language::Language>,
879 ) -> Option<language::CodeLabel> {
880 let label = &item.label;
881 let grammar = language.grammar()?;
882 let highlight_id = match item.kind? {
883 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
884 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
885 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
886 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
887 _ => return None,
888 };
889 Some(language::CodeLabel {
890 text: label.clone(),
891 runs: vec![(0..label.len(), highlight_id)],
892 filter_range: 0..label.len(),
893 })
894 }
895
896 async fn label_for_symbol(
897 &self,
898 name: &str,
899 kind: lsp::SymbolKind,
900 language: &Arc<language::Language>,
901 ) -> Option<language::CodeLabel> {
902 let (text, filter_range, display_range) = match kind {
903 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
904 let text = format!("def {}():\n", name);
905 let filter_range = 4..4 + name.len();
906 let display_range = 0..filter_range.end;
907 (text, filter_range, display_range)
908 }
909 lsp::SymbolKind::CLASS => {
910 let text = format!("class {}:", name);
911 let filter_range = 6..6 + name.len();
912 let display_range = 0..filter_range.end;
913 (text, filter_range, display_range)
914 }
915 lsp::SymbolKind::CONSTANT => {
916 let text = format!("{} = 0", name);
917 let filter_range = 0..name.len();
918 let display_range = 0..filter_range.end;
919 (text, filter_range, display_range)
920 }
921 _ => return None,
922 };
923
924 Some(language::CodeLabel {
925 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
926 text: text[display_range].to_string(),
927 filter_range,
928 })
929 }
930
931 async fn workspace_configuration(
932 self: Arc<Self>,
933 adapter: &Arc<dyn LspAdapterDelegate>,
934 toolchains: Arc<dyn LanguageToolchainStore>,
935 cx: &mut AsyncAppContext,
936 ) -> Result<Value> {
937 let toolchain = toolchains
938 .active_toolchain(adapter.worktree_id(), LanguageName::new("Python"), cx)
939 .await;
940 cx.update(move |cx| {
941 let mut user_settings =
942 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
943 .and_then(|s| s.settings.clone())
944 .unwrap_or_else(|| {
945 json!({
946 "plugins": {
947 "pycodestyle": {"enabled": false},
948 "rope_autoimport": {"enabled": true, "memory": true},
949 "pylsp_mypy": {"enabled": false}
950 },
951 "rope": {
952 "ropeFolder": null
953 },
954 })
955 });
956
957 // If user did not explicitly modify their python venv, use one from picker.
958 if let Some(toolchain) = toolchain {
959 if user_settings.is_null() {
960 user_settings = Value::Object(serde_json::Map::default());
961 }
962 let object = user_settings.as_object_mut().unwrap();
963 if let Some(python) = object
964 .entry("plugins")
965 .or_insert(Value::Object(serde_json::Map::default()))
966 .as_object_mut()
967 {
968 if let Some(jedi) = python
969 .entry("jedi")
970 .or_insert(Value::Object(serde_json::Map::default()))
971 .as_object_mut()
972 {
973 jedi.entry("environment".to_string())
974 .or_insert_with(|| Value::String(toolchain.path.clone().into()));
975 }
976 if let Some(pylint) = python
977 .entry("pylsp_mypy")
978 .or_insert(Value::Object(serde_json::Map::default()))
979 .as_object_mut()
980 {
981 pylint.entry("overrides".to_string()).or_insert_with(|| {
982 Value::Array(vec![
983 Value::String("--python-executable".into()),
984 Value::String(toolchain.path.into()),
985 Value::String("--cache-dir=/dev/null".into()),
986 Value::Bool(true),
987 ])
988 });
989 }
990 }
991 }
992 user_settings = Value::Object(serde_json::Map::from_iter([(
993 "pylsp".to_string(),
994 user_settings,
995 )]));
996
997 user_settings
998 })
999 }
1000}
1001
1002#[cfg(test)]
1003mod tests {
1004 use gpui::{BorrowAppContext, Context, ModelContext, TestAppContext};
1005 use language::{language_settings::AllLanguageSettings, AutoindentMode, Buffer};
1006 use settings::SettingsStore;
1007 use std::num::NonZeroU32;
1008
1009 #[gpui::test]
1010 async fn test_python_autoindent(cx: &mut TestAppContext) {
1011 cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
1012 let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
1013 cx.update(|cx| {
1014 let test_settings = SettingsStore::test(cx);
1015 cx.set_global(test_settings);
1016 language::init(cx);
1017 cx.update_global::<SettingsStore, _>(|store, cx| {
1018 store.update_user_settings::<AllLanguageSettings>(cx, |s| {
1019 s.defaults.tab_size = NonZeroU32::new(2);
1020 });
1021 });
1022 });
1023
1024 cx.new_model(|cx| {
1025 let mut buffer = Buffer::local("", cx).with_language(language, cx);
1026 let append = |buffer: &mut Buffer, text: &str, cx: &mut ModelContext<Buffer>| {
1027 let ix = buffer.len();
1028 buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
1029 };
1030
1031 // indent after "def():"
1032 append(&mut buffer, "def a():\n", cx);
1033 assert_eq!(buffer.text(), "def a():\n ");
1034
1035 // preserve indent after blank line
1036 append(&mut buffer, "\n ", cx);
1037 assert_eq!(buffer.text(), "def a():\n \n ");
1038
1039 // indent after "if"
1040 append(&mut buffer, "if a:\n ", cx);
1041 assert_eq!(buffer.text(), "def a():\n \n if a:\n ");
1042
1043 // preserve indent after statement
1044 append(&mut buffer, "b()\n", cx);
1045 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n ");
1046
1047 // preserve indent after statement
1048 append(&mut buffer, "else", cx);
1049 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else");
1050
1051 // dedent "else""
1052 append(&mut buffer, ":", cx);
1053 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else:");
1054
1055 // indent lines after else
1056 append(&mut buffer, "\n", cx);
1057 assert_eq!(
1058 buffer.text(),
1059 "def a():\n \n if a:\n b()\n else:\n "
1060 );
1061
1062 // indent after an open paren. the closing paren is not indented
1063 // because there is another token before it on the same line.
1064 append(&mut buffer, "foo(\n1)", cx);
1065 assert_eq!(
1066 buffer.text(),
1067 "def a():\n \n if a:\n b()\n else:\n foo(\n 1)"
1068 );
1069
1070 // dedent the closing paren if it is shifted to the beginning of the line
1071 let argument_ix = buffer.text().find('1').unwrap();
1072 buffer.edit(
1073 [(argument_ix..argument_ix + 1, "")],
1074 Some(AutoindentMode::EachLine),
1075 cx,
1076 );
1077 assert_eq!(
1078 buffer.text(),
1079 "def a():\n \n if a:\n b()\n else:\n foo(\n )"
1080 );
1081
1082 // preserve indent after the close paren
1083 append(&mut buffer, "\n", cx);
1084 assert_eq!(
1085 buffer.text(),
1086 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n "
1087 );
1088
1089 // manually outdent the last line
1090 let end_whitespace_ix = buffer.len() - 4;
1091 buffer.edit(
1092 [(end_whitespace_ix..buffer.len(), "")],
1093 Some(AutoindentMode::EachLine),
1094 cx,
1095 );
1096 assert_eq!(
1097 buffer.text(),
1098 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n"
1099 );
1100
1101 // preserve the newly reduced indentation on the next newline
1102 append(&mut buffer, "\n", cx);
1103 assert_eq!(
1104 buffer.text(),
1105 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n\n"
1106 );
1107
1108 // reset to a simple if statement
1109 buffer.edit([(0..buffer.len(), "if a:\n b(\n )")], None, cx);
1110
1111 // dedent "else" on the line after a closing paren
1112 append(&mut buffer, "\n else:\n", cx);
1113 assert_eq!(buffer.text(), "if a:\n b(\n )\nelse:\n ");
1114
1115 buffer
1116 });
1117 }
1118}