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