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