1use anyhow::Result;
2use async_trait::async_trait;
3use collections::HashMap;
4use gpui::AppContext;
5use gpui::AsyncAppContext;
6use language::LanguageName;
7use language::LanguageToolchainStore;
8use language::Toolchain;
9use language::ToolchainList;
10use language::ToolchainLister;
11use language::{ContextProvider, LspAdapter, LspAdapterDelegate};
12use lsp::LanguageServerBinary;
13use lsp::LanguageServerName;
14use node_runtime::NodeRuntime;
15use pet_core::os_environment::Environment;
16use pet_core::python_environment::PythonEnvironmentKind;
17use pet_core::Configuration;
18use project::lsp_store::language_server_settings;
19use serde_json::Value;
20
21use std::sync::Mutex;
22use std::{
23 any::Any,
24 borrow::Cow,
25 ffi::OsString,
26 path::{Path, PathBuf},
27 sync::Arc,
28};
29use task::{TaskTemplate, TaskTemplates, VariableName};
30use util::ResultExt;
31
32const SERVER_PATH: &str = "node_modules/pyright/langserver.index.js";
33const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "pyright/langserver.index.js";
34
35fn server_binary_arguments(server_path: &Path) -> Vec<OsString> {
36 vec![server_path.into(), "--stdio".into()]
37}
38
39pub struct PythonLspAdapter {
40 node: NodeRuntime,
41}
42
43impl PythonLspAdapter {
44 const SERVER_NAME: LanguageServerName = LanguageServerName::new_static("pyright");
45
46 pub fn new(node: NodeRuntime) -> Self {
47 PythonLspAdapter { node }
48 }
49}
50
51#[async_trait(?Send)]
52impl LspAdapter for PythonLspAdapter {
53 fn name(&self) -> LanguageServerName {
54 Self::SERVER_NAME.clone()
55 }
56
57 async fn check_if_user_installed(
58 &self,
59 delegate: &dyn LspAdapterDelegate,
60 _: &AsyncAppContext,
61 ) -> Option<LanguageServerBinary> {
62 let node = delegate.which("node".as_ref()).await?;
63 let (node_modules_path, _) = delegate
64 .npm_package_installed_version(Self::SERVER_NAME.as_ref())
65 .await
66 .log_err()??;
67
68 let path = node_modules_path.join(NODE_MODULE_RELATIVE_SERVER_PATH);
69
70 Some(LanguageServerBinary {
71 path: node,
72 env: None,
73 arguments: server_binary_arguments(&path),
74 })
75 }
76
77 async fn fetch_latest_server_version(
78 &self,
79 _: &dyn LspAdapterDelegate,
80 ) -> Result<Box<dyn 'static + Any + Send>> {
81 Ok(Box::new(
82 self.node
83 .npm_package_latest_version(Self::SERVER_NAME.as_ref())
84 .await?,
85 ) as Box<_>)
86 }
87
88 async fn fetch_server_binary(
89 &self,
90 latest_version: Box<dyn 'static + Send + Any>,
91 container_dir: PathBuf,
92 _: &dyn LspAdapterDelegate,
93 ) -> Result<LanguageServerBinary> {
94 let latest_version = latest_version.downcast::<String>().unwrap();
95 let server_path = container_dir.join(SERVER_PATH);
96
97 let should_install_language_server = self
98 .node
99 .should_install_npm_package(
100 Self::SERVER_NAME.as_ref(),
101 &server_path,
102 &container_dir,
103 &latest_version,
104 )
105 .await;
106
107 if should_install_language_server {
108 self.node
109 .npm_install_packages(
110 &container_dir,
111 &[(Self::SERVER_NAME.as_ref(), latest_version.as_str())],
112 )
113 .await?;
114 }
115
116 Ok(LanguageServerBinary {
117 path: self.node.binary_path().await?,
118 env: None,
119 arguments: server_binary_arguments(&server_path),
120 })
121 }
122
123 async fn cached_server_binary(
124 &self,
125 container_dir: PathBuf,
126 _: &dyn LspAdapterDelegate,
127 ) -> Option<LanguageServerBinary> {
128 get_cached_server_binary(container_dir, &self.node).await
129 }
130
131 async fn process_completions(&self, items: &mut [lsp::CompletionItem]) {
132 // Pyright assigns each completion item a `sortText` of the form `XX.YYYY.name`.
133 // Where `XX` is the sorting category, `YYYY` is based on most recent usage,
134 // and `name` is the symbol name itself.
135 //
136 // Because the symbol name is included, there generally are not ties when
137 // sorting by the `sortText`, so the symbol's fuzzy match score is not taken
138 // into account. Here, we remove the symbol name from the sortText in order
139 // to allow our own fuzzy score to be used to break ties.
140 //
141 // see https://github.com/microsoft/pyright/blob/95ef4e103b9b2f129c9320427e51b73ea7cf78bd/packages/pyright-internal/src/languageService/completionProvider.ts#LL2873
142 for item in items {
143 let Some(sort_text) = &mut item.sort_text else {
144 continue;
145 };
146 let mut parts = sort_text.split('.');
147 let Some(first) = parts.next() else { continue };
148 let Some(second) = parts.next() else { continue };
149 let Some(_) = parts.next() else { continue };
150 sort_text.replace_range(first.len() + second.len() + 1.., "");
151 }
152 }
153
154 async fn label_for_completion(
155 &self,
156 item: &lsp::CompletionItem,
157 language: &Arc<language::Language>,
158 ) -> Option<language::CodeLabel> {
159 let label = &item.label;
160 let grammar = language.grammar()?;
161 let highlight_id = match item.kind? {
162 lsp::CompletionItemKind::METHOD => grammar.highlight_id_for_name("function.method")?,
163 lsp::CompletionItemKind::FUNCTION => grammar.highlight_id_for_name("function")?,
164 lsp::CompletionItemKind::CLASS => grammar.highlight_id_for_name("type")?,
165 lsp::CompletionItemKind::CONSTANT => grammar.highlight_id_for_name("constant")?,
166 _ => return None,
167 };
168 Some(language::CodeLabel {
169 text: label.clone(),
170 runs: vec![(0..label.len(), highlight_id)],
171 filter_range: 0..label.len(),
172 })
173 }
174
175 async fn label_for_symbol(
176 &self,
177 name: &str,
178 kind: lsp::SymbolKind,
179 language: &Arc<language::Language>,
180 ) -> Option<language::CodeLabel> {
181 let (text, filter_range, display_range) = match kind {
182 lsp::SymbolKind::METHOD | lsp::SymbolKind::FUNCTION => {
183 let text = format!("def {}():\n", name);
184 let filter_range = 4..4 + name.len();
185 let display_range = 0..filter_range.end;
186 (text, filter_range, display_range)
187 }
188 lsp::SymbolKind::CLASS => {
189 let text = format!("class {}:", name);
190 let filter_range = 6..6 + name.len();
191 let display_range = 0..filter_range.end;
192 (text, filter_range, display_range)
193 }
194 lsp::SymbolKind::CONSTANT => {
195 let text = format!("{} = 0", name);
196 let filter_range = 0..name.len();
197 let display_range = 0..filter_range.end;
198 (text, filter_range, display_range)
199 }
200 _ => return None,
201 };
202
203 Some(language::CodeLabel {
204 runs: language.highlight_text(&text.as_str().into(), display_range.clone()),
205 text: text[display_range].to_string(),
206 filter_range,
207 })
208 }
209
210 async fn workspace_configuration(
211 self: Arc<Self>,
212 adapter: &Arc<dyn LspAdapterDelegate>,
213 toolchains: Arc<dyn LanguageToolchainStore>,
214 cx: &mut AsyncAppContext,
215 ) -> Result<Value> {
216 let toolchain = toolchains
217 .active_toolchain(adapter.worktree_id(), LanguageName::new("Python"), cx)
218 .await;
219 cx.update(move |cx| {
220 let mut user_settings =
221 language_server_settings(adapter.as_ref(), &Self::SERVER_NAME, cx)
222 .and_then(|s| s.settings.clone())
223 .unwrap_or_default();
224
225 // If python.pythonPath is not set in user config, do so using our toolchain picker.
226 if let Some(toolchain) = toolchain {
227 if user_settings.is_null() {
228 user_settings = Value::Object(serde_json::Map::default());
229 }
230 let object = user_settings.as_object_mut().unwrap();
231 if let Some(python) = object
232 .entry("python")
233 .or_insert(Value::Object(serde_json::Map::default()))
234 .as_object_mut()
235 {
236 python
237 .entry("pythonPath")
238 .or_insert(Value::String(toolchain.path.into()));
239 }
240 }
241 user_settings
242 })
243 }
244}
245
246async fn get_cached_server_binary(
247 container_dir: PathBuf,
248 node: &NodeRuntime,
249) -> Option<LanguageServerBinary> {
250 let server_path = container_dir.join(SERVER_PATH);
251 if server_path.exists() {
252 Some(LanguageServerBinary {
253 path: node.binary_path().await.log_err()?,
254 env: None,
255 arguments: server_binary_arguments(&server_path),
256 })
257 } else {
258 log::error!("missing executable in directory {:?}", server_path);
259 None
260 }
261}
262
263pub(crate) struct PythonContextProvider;
264
265const PYTHON_UNITTEST_TARGET_TASK_VARIABLE: VariableName =
266 VariableName::Custom(Cow::Borrowed("PYTHON_UNITTEST_TARGET"));
267
268impl ContextProvider for PythonContextProvider {
269 fn build_context(
270 &self,
271 variables: &task::TaskVariables,
272 _location: &project::Location,
273 _: Option<&HashMap<String, String>>,
274 _cx: &mut gpui::AppContext,
275 ) -> Result<task::TaskVariables> {
276 let python_module_name = python_module_name_from_relative_path(
277 variables.get(&VariableName::RelativeFile).unwrap_or(""),
278 );
279 let unittest_class_name =
280 variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
281 let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
282 "_unittest_method_name",
283 )));
284
285 let unittest_target_str = match (unittest_class_name, unittest_method_name) {
286 (Some(class_name), Some(method_name)) => {
287 format!("{}.{}.{}", python_module_name, class_name, method_name)
288 }
289 (Some(class_name), None) => format!("{}.{}", python_module_name, class_name),
290 (None, None) => python_module_name,
291 (None, Some(_)) => return Ok(task::TaskVariables::default()), // should never happen, a TestCase class is the unit of testing
292 };
293
294 let unittest_target = (
295 PYTHON_UNITTEST_TARGET_TASK_VARIABLE.clone(),
296 unittest_target_str,
297 );
298
299 Ok(task::TaskVariables::from_iter([unittest_target]))
300 }
301
302 fn associated_tasks(
303 &self,
304 _: Option<Arc<dyn language::File>>,
305 _: &AppContext,
306 ) -> Option<TaskTemplates> {
307 Some(TaskTemplates(vec![
308 TaskTemplate {
309 label: "execute selection".to_owned(),
310 command: "python3".to_owned(),
311 args: vec!["-c".to_owned(), VariableName::SelectedText.template_value()],
312 ..TaskTemplate::default()
313 },
314 TaskTemplate {
315 label: format!("run '{}'", VariableName::File.template_value()),
316 command: "python3".to_owned(),
317 args: vec![VariableName::File.template_value()],
318 ..TaskTemplate::default()
319 },
320 TaskTemplate {
321 label: format!("unittest '{}'", VariableName::File.template_value()),
322 command: "python3".to_owned(),
323 args: vec![
324 "-m".to_owned(),
325 "unittest".to_owned(),
326 VariableName::File.template_value(),
327 ],
328 ..TaskTemplate::default()
329 },
330 TaskTemplate {
331 label: "unittest $ZED_CUSTOM_PYTHON_UNITTEST_TARGET".to_owned(),
332 command: "python3".to_owned(),
333 args: vec![
334 "-m".to_owned(),
335 "unittest".to_owned(),
336 "$ZED_CUSTOM_PYTHON_UNITTEST_TARGET".to_owned(),
337 ],
338 tags: vec![
339 "python-unittest-class".to_owned(),
340 "python-unittest-method".to_owned(),
341 ],
342 ..TaskTemplate::default()
343 },
344 ]))
345 }
346}
347
348fn python_module_name_from_relative_path(relative_path: &str) -> String {
349 let path_with_dots = relative_path.replace('/', ".");
350 path_with_dots
351 .strip_suffix(".py")
352 .unwrap_or(&path_with_dots)
353 .to_string()
354}
355
356#[derive(Default)]
357pub(crate) struct PythonToolchainProvider {}
358
359static ENV_PRIORITY_LIST: &'static [PythonEnvironmentKind] = &[
360 // Prioritize non-Conda environments.
361 PythonEnvironmentKind::Poetry,
362 PythonEnvironmentKind::Pipenv,
363 PythonEnvironmentKind::VirtualEnvWrapper,
364 PythonEnvironmentKind::Venv,
365 PythonEnvironmentKind::VirtualEnv,
366 PythonEnvironmentKind::Conda,
367 PythonEnvironmentKind::Pyenv,
368 PythonEnvironmentKind::GlobalPaths,
369 PythonEnvironmentKind::Homebrew,
370];
371
372fn env_priority(kind: Option<PythonEnvironmentKind>) -> usize {
373 if let Some(kind) = kind {
374 ENV_PRIORITY_LIST
375 .iter()
376 .position(|blessed_env| blessed_env == &kind)
377 .unwrap_or(ENV_PRIORITY_LIST.len())
378 } else {
379 // Unknown toolchains are less useful than non-blessed ones.
380 ENV_PRIORITY_LIST.len() + 1
381 }
382}
383
384#[async_trait(?Send)]
385impl ToolchainLister for PythonToolchainProvider {
386 async fn list(
387 &self,
388 worktree_root: PathBuf,
389 project_env: Option<HashMap<String, String>>,
390 ) -> ToolchainList {
391 let env = project_env.unwrap_or_default();
392 let environment = EnvironmentApi::from_env(&env);
393 let locators = pet::locators::create_locators(
394 Arc::new(pet_conda::Conda::from(&environment)),
395 Arc::new(pet_poetry::Poetry::from(&environment)),
396 &environment,
397 );
398 let mut config = Configuration::default();
399 config.workspace_directories = Some(vec![worktree_root]);
400 let reporter = pet_reporter::collect::create_reporter();
401 pet::find::find_and_report_envs(&reporter, config, &locators, &environment, None);
402
403 let mut toolchains = reporter
404 .environments
405 .lock()
406 .ok()
407 .map_or(Vec::new(), |mut guard| std::mem::take(&mut guard));
408 toolchains.sort_by(|lhs, rhs| {
409 env_priority(lhs.kind)
410 .cmp(&env_priority(rhs.kind))
411 .then_with(|| lhs.executable.cmp(&rhs.executable))
412 });
413 let mut toolchains: Vec<_> = toolchains
414 .into_iter()
415 .filter_map(|toolchain| {
416 let name = if let Some(version) = &toolchain.version {
417 format!("Python {version} ({:?})", toolchain.kind?)
418 } else {
419 format!("{:?}", toolchain.kind?)
420 }
421 .into();
422 Some(Toolchain {
423 name,
424 path: toolchain.executable?.to_str()?.to_owned().into(),
425 language_name: LanguageName::new("Python"),
426 })
427 })
428 .collect();
429 toolchains.dedup();
430 ToolchainList {
431 toolchains,
432 default: None,
433 groups: Default::default(),
434 }
435 }
436}
437
438pub struct EnvironmentApi<'a> {
439 global_search_locations: Arc<Mutex<Vec<PathBuf>>>,
440 project_env: &'a HashMap<String, String>,
441 pet_env: pet_core::os_environment::EnvironmentApi,
442}
443
444impl<'a> EnvironmentApi<'a> {
445 pub fn from_env(project_env: &'a HashMap<String, String>) -> Self {
446 let paths = project_env
447 .get("PATH")
448 .map(|p| std::env::split_paths(p).collect())
449 .unwrap_or_default();
450
451 EnvironmentApi {
452 global_search_locations: Arc::new(Mutex::new(paths)),
453 project_env,
454 pet_env: pet_core::os_environment::EnvironmentApi::new(),
455 }
456 }
457
458 fn user_home(&self) -> Option<PathBuf> {
459 self.project_env
460 .get("HOME")
461 .or_else(|| self.project_env.get("USERPROFILE"))
462 .map(|home| pet_fs::path::norm_case(PathBuf::from(home)))
463 .or_else(|| self.pet_env.get_user_home())
464 }
465}
466
467impl<'a> pet_core::os_environment::Environment for EnvironmentApi<'a> {
468 fn get_user_home(&self) -> Option<PathBuf> {
469 self.user_home()
470 }
471
472 fn get_root(&self) -> Option<PathBuf> {
473 None
474 }
475
476 fn get_env_var(&self, key: String) -> Option<String> {
477 self.project_env
478 .get(&key)
479 .cloned()
480 .or_else(|| self.pet_env.get_env_var(key))
481 }
482
483 fn get_know_global_search_locations(&self) -> Vec<PathBuf> {
484 if self.global_search_locations.lock().unwrap().is_empty() {
485 let mut paths =
486 std::env::split_paths(&self.get_env_var("PATH".to_string()).unwrap_or_default())
487 .collect::<Vec<PathBuf>>();
488
489 log::trace!("Env PATH: {:?}", paths);
490 for p in self.pet_env.get_know_global_search_locations() {
491 if !paths.contains(&p) {
492 paths.push(p);
493 }
494 }
495
496 let mut paths = paths
497 .into_iter()
498 .filter(|p| p.exists())
499 .collect::<Vec<PathBuf>>();
500
501 self.global_search_locations
502 .lock()
503 .unwrap()
504 .append(&mut paths);
505 }
506 self.global_search_locations.lock().unwrap().clone()
507 }
508}
509
510#[cfg(test)]
511mod tests {
512 use gpui::{BorrowAppContext, Context, ModelContext, TestAppContext};
513 use language::{language_settings::AllLanguageSettings, AutoindentMode, Buffer};
514 use settings::SettingsStore;
515 use std::num::NonZeroU32;
516
517 #[gpui::test]
518 async fn test_python_autoindent(cx: &mut TestAppContext) {
519 cx.executor().set_block_on_ticks(usize::MAX..=usize::MAX);
520 let language = crate::language("python", tree_sitter_python::LANGUAGE.into());
521 cx.update(|cx| {
522 let test_settings = SettingsStore::test(cx);
523 cx.set_global(test_settings);
524 language::init(cx);
525 cx.update_global::<SettingsStore, _>(|store, cx| {
526 store.update_user_settings::<AllLanguageSettings>(cx, |s| {
527 s.defaults.tab_size = NonZeroU32::new(2);
528 });
529 });
530 });
531
532 cx.new_model(|cx| {
533 let mut buffer = Buffer::local("", cx).with_language(language, cx);
534 let append = |buffer: &mut Buffer, text: &str, cx: &mut ModelContext<Buffer>| {
535 let ix = buffer.len();
536 buffer.edit([(ix..ix, text)], Some(AutoindentMode::EachLine), cx);
537 };
538
539 // indent after "def():"
540 append(&mut buffer, "def a():\n", cx);
541 assert_eq!(buffer.text(), "def a():\n ");
542
543 // preserve indent after blank line
544 append(&mut buffer, "\n ", cx);
545 assert_eq!(buffer.text(), "def a():\n \n ");
546
547 // indent after "if"
548 append(&mut buffer, "if a:\n ", cx);
549 assert_eq!(buffer.text(), "def a():\n \n if a:\n ");
550
551 // preserve indent after statement
552 append(&mut buffer, "b()\n", cx);
553 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n ");
554
555 // preserve indent after statement
556 append(&mut buffer, "else", cx);
557 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else");
558
559 // dedent "else""
560 append(&mut buffer, ":", cx);
561 assert_eq!(buffer.text(), "def a():\n \n if a:\n b()\n else:");
562
563 // indent lines after else
564 append(&mut buffer, "\n", cx);
565 assert_eq!(
566 buffer.text(),
567 "def a():\n \n if a:\n b()\n else:\n "
568 );
569
570 // indent after an open paren. the closing paren is not indented
571 // because there is another token before it on the same line.
572 append(&mut buffer, "foo(\n1)", cx);
573 assert_eq!(
574 buffer.text(),
575 "def a():\n \n if a:\n b()\n else:\n foo(\n 1)"
576 );
577
578 // dedent the closing paren if it is shifted to the beginning of the line
579 let argument_ix = buffer.text().find('1').unwrap();
580 buffer.edit(
581 [(argument_ix..argument_ix + 1, "")],
582 Some(AutoindentMode::EachLine),
583 cx,
584 );
585 assert_eq!(
586 buffer.text(),
587 "def a():\n \n if a:\n b()\n else:\n foo(\n )"
588 );
589
590 // preserve indent after the close paren
591 append(&mut buffer, "\n", cx);
592 assert_eq!(
593 buffer.text(),
594 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n "
595 );
596
597 // manually outdent the last line
598 let end_whitespace_ix = buffer.len() - 4;
599 buffer.edit(
600 [(end_whitespace_ix..buffer.len(), "")],
601 Some(AutoindentMode::EachLine),
602 cx,
603 );
604 assert_eq!(
605 buffer.text(),
606 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n"
607 );
608
609 // preserve the newly reduced indentation on the next newline
610 append(&mut buffer, "\n", cx);
611 assert_eq!(
612 buffer.text(),
613 "def a():\n \n if a:\n b()\n else:\n foo(\n )\n\n"
614 );
615
616 // reset to a simple if statement
617 buffer.edit([(0..buffer.len(), "if a:\n b(\n )")], None, cx);
618
619 // dedent "else" on the line after a closing paren
620 append(&mut buffer, "\n else:\n", cx);
621 assert_eq!(buffer.text(), "if a:\n b(\n )\nelse:\n ");
622
623 buffer
624 });
625 }
626}