From 1574a3a2fdc6bfbb4493dc3a1ee1da9dbed807ab Mon Sep 17 00:00:00 2001
From: Alex van de Griendt <avdg.lucerna@gmail.com>
Date: Tue, 25 Mar 2025 22:08:16 +0100
Subject: [PATCH] python: Add task for running modules (#26462)

Closes #26460

I am new to contributing to Zed (and pretty new to Rust in general). I'm
not too familiar with code style, guidelines etc. so please feel free to
suggest changes/improvements.

This PR adds a run icon to Python files that have a "main" function:
```python
if __name__ == "__main__":
    ...
```

In addition to the gutter icon, there is now also an extra task in the
command palette "run module".

Release Notes:

- Added detection for runnable Python modules
- Added Python-specific task to run a Python file as a module from
inside the project's scope

---------

Co-authored-by: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
---
 crates/languages/src/python.rs            | 36 ++++++++++++++++++++++-
 crates/languages/src/python/runnables.scm | 16 ++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/crates/languages/src/python.rs b/crates/languages/src/python.rs
index 825b11db2d..250b80cc32 100644
--- a/crates/languages/src/python.rs
+++ b/crates/languages/src/python.rs
@@ -348,6 +348,10 @@ const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
 
 const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
     VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
+
+const PYTHON_MODULE_NAME_TASK_VARIABLE: VariableName =
+    VariableName::Custom(Cow::Borrowed("PYTHON_MODULE_NAME"));
+
 impl ContextProvider for PythonContextProvider {
     fn build_context(
         &self,
@@ -362,7 +366,9 @@ impl ContextProvider for PythonContextProvider {
             TestRunner::PYTEST => self.build_pytest_target(variables),
         };
 
+        let module_target = self.build_module_target(variables);
         let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
+
         cx.spawn(async move |cx| {
             let active_toolchain = if let Some(worktree_id) = worktree_id {
                 toolchains
@@ -376,8 +382,12 @@ impl ContextProvider for PythonContextProvider {
                 String::from("python3")
             };
             let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
+
             Ok(task::TaskVariables::from_iter(
-                test_target.into_iter().chain([toolchain]),
+                test_target
+                    .into_iter()
+                    .chain(module_target.into_iter())
+                    .chain([toolchain]),
             ))
         })
     }
@@ -407,6 +417,17 @@ impl ContextProvider for PythonContextProvider {
                 args: vec![VariableName::File.template_value_with_whitespace()],
                 ..TaskTemplate::default()
             },
+            // Execute a file as module
+            TaskTemplate {
+                label: format!("run module '{}'", VariableName::File.template_value()),
+                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
+                args: vec![
+                    "-m".to_owned(),
+                    PYTHON_MODULE_NAME_TASK_VARIABLE.template_value(),
+                ],
+                tags: vec!["python-module-main-method".to_owned()],
+                ..TaskTemplate::default()
+            },
         ];
 
         tasks.extend(match test_runner {
@@ -544,6 +565,19 @@ impl PythonContextProvider {
 
         Some((PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str))
     }
+
+    fn build_module_target(
+        &self,
+        variables: &task::TaskVariables,
+    ) -> Result<(VariableName, String)> {
+        let python_module_name = python_module_name_from_relative_path(
+            variables.get(&VariableName::RelativeFile).unwrap_or(""),
+        );
+
+        let module_target = (PYTHON_MODULE_NAME_TASK_VARIABLE.clone(), python_module_name);
+
+        Ok(module_target)
+    }
 }
 
 fn python_module_name_from_relative_path(relative_path: &str) -> String {
diff --git a/crates/languages/src/python/runnables.scm b/crates/languages/src/python/runnables.scm
index 8cdb0d77eb..3b32556707 100644
--- a/crates/languages/src/python/runnables.scm
+++ b/crates/languages/src/python/runnables.scm
@@ -82,3 +82,19 @@
         )
     )
 )
+
+; module main method
+(
+    (module
+        (if_statement
+            condition: (comparison_operator
+                (identifier) @run @_lhs
+                operators: "=="
+                (string) @_rhs
+            )
+            (#eq? @_lhs "__name__")
+            (#match? @_rhs "^[\"']__main__[\"']$")
+            (#set! tag python-module-main-method)
+        )
+    )
+)