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