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