Add pytest-based test discovery and runnables for Python (#18824)

Closes  #12080, #18649.

Screenshot:

<img width="1499" alt="image"
src="https://github.com/user-attachments/assets/2644c2fc-19cf-4d2c-a992-5c56cb22deed">

Still in progress:

1. I'd like to add configuration options for selecting a Python test
runner (either pytest or unittest) so that users can explicitly choose
which runner they'd like to use for running their tests. This preference
has to be configured as unittest-style tests can also be run by pytest,
meaning we can't rely on auto-discovery to choose the desired test
runner.
2. I'd like to add venv auto-discovery similar to the feature currently
provided by the terminal using detect_venv.
3. Unit tests.

Unfortunately I'm struggling a bit with how to add settings in the
appropriate location (e.g. Python language settings). Can anyone provide
me with some pointers and/or examples on how to either add extra
settings or to re-use the existing ones?

My rust programming level is OK-ish but I'm not very familiar with the
Zed project structure and could use some help.

I'm also open for pair programming as mentioned on the website if that
helps!

Release Notes:

- Added pytest-based test discovery and runnables for Python.
- Adds a configurable option for switching between unittest and pytest
as a test runner under Python language settings. Set "TASK_RUNNER" to
"unittest" under task settings for Python if you wish to use unittest to
run Python tasks; the default is pytest.

---------

Co-authored-by: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>

crates/languages/src/python.rs

@@ -4,6 +4,7 @@ use async_trait::async_trait;
 use collections::HashMap;
 use gpui::AsyncAppContext;
 use gpui::{AppContext, Task};
+use language::language_settings::language_settings;
 use language::LanguageName;
 use language::LanguageToolchainStore;
 use language::Toolchain;
@@ -21,6 +22,7 @@ use serde_json::{json, Value};
 use smol::{lock::OnceCell, process::Command};
 use std::cmp::Ordering;
 
+use std::str::FromStr;
 use std::sync::Mutex;
 use std::{
     any::Any,
@@ -35,6 +37,23 @@ use util::ResultExt;
 const SERVER_PATH: &str = "node_modules/pyright/langserver.index.js";
 const NODE_MODULE_RELATIVE_SERVER_PATH: &str = "pyright/langserver.index.js";
 
+enum TestRunner {
+    UNITTEST,
+    PYTEST,
+}
+
+impl FromStr for TestRunner {
+    type Err = ();
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s {
+            "unittest" => Ok(Self::UNITTEST),
+            "pytest" => Ok(Self::PYTEST),
+            _ => Err(()),
+        }
+    }
+}
+
 fn server_binary_arguments(server_path: &Path) -> Vec<OsString> {
     vec![server_path.into(), "--stdio".into()]
 }
@@ -265,8 +284,8 @@ async fn get_cached_server_binary(
 
 pub(crate) struct PythonContextProvider;
 
-const PYTHON_UNITTEST_TARGET_TASK_VARIABLE: VariableName =
-    VariableName::Custom(Cow::Borrowed("PYTHON_UNITTEST_TARGET"));
+const PYTHON_TEST_TARGET_TASK_VARIABLE: VariableName =
+    VariableName::Custom(Cow::Borrowed("PYTHON_TEST_TARGET"));
 
 const PYTHON_ACTIVE_TOOLCHAIN_PATH: VariableName =
     VariableName::Custom(Cow::Borrowed("PYTHON_ACTIVE_ZED_TOOLCHAIN"));
@@ -279,28 +298,16 @@ impl ContextProvider for PythonContextProvider {
         toolchains: Arc<dyn LanguageToolchainStore>,
         cx: &mut gpui::AppContext,
     ) -> Task<Result<task::TaskVariables>> {
-        let python_module_name = python_module_name_from_relative_path(
-            variables.get(&VariableName::RelativeFile).unwrap_or(""),
-        );
-        let unittest_class_name =
-            variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
-        let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
-            "_unittest_method_name",
-        )));
+        let test_target = {
+            let test_runner = selected_test_runner(location.buffer.read(cx).file(), cx);
 
-        let unittest_target_str = match (unittest_class_name, unittest_method_name) {
-            (Some(class_name), Some(method_name)) => {
-                format!("{}.{}.{}", python_module_name, class_name, method_name)
-            }
-            (Some(class_name), None) => format!("{}.{}", python_module_name, class_name),
-            (None, None) => python_module_name,
-            (None, Some(_)) => return Task::ready(Ok(task::TaskVariables::default())), // should never happen, a TestCase class is the unit of testing
+            let runner = match test_runner {
+                TestRunner::UNITTEST => self.build_unittest_target(variables),
+                TestRunner::PYTEST => self.build_pytest_target(variables),
+            };
+            runner
         };
 
-        let unittest_target = (
-            PYTHON_UNITTEST_TARGET_TASK_VARIABLE.clone(),
-            unittest_target_str,
-        );
         let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
         cx.spawn(move |mut cx| async move {
             let active_toolchain = if let Some(worktree_id) = worktree_id {
@@ -312,53 +319,174 @@ impl ContextProvider for PythonContextProvider {
                 String::from("python3")
             };
             let toolchain = (PYTHON_ACTIVE_TOOLCHAIN_PATH, active_toolchain);
-            Ok(task::TaskVariables::from_iter([unittest_target, toolchain]))
+            Ok(task::TaskVariables::from_iter([test_target?, toolchain]))
         })
     }
 
     fn associated_tasks(
         &self,
-        _: Option<Arc<dyn language::File>>,
-        _: &AppContext,
+        file: Option<Arc<dyn language::File>>,
+        cx: &AppContext,
     ) -> Option<TaskTemplates> {
-        Some(TaskTemplates(vec![
+        let test_runner = selected_test_runner(file.as_ref(), cx);
+
+        let mut tasks = vec![
+            // Execute a selection
             TaskTemplate {
                 label: "execute selection".to_owned(),
                 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
                 args: vec!["-c".to_owned(), VariableName::SelectedText.template_value()],
                 ..TaskTemplate::default()
             },
+            // Execute an entire file
             TaskTemplate {
                 label: format!("run '{}'", VariableName::File.template_value()),
                 command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
                 args: vec![VariableName::File.template_value()],
                 ..TaskTemplate::default()
             },
-            TaskTemplate {
-                label: format!("unittest '{}'", VariableName::File.template_value()),
-                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
-                args: vec![
-                    "-m".to_owned(),
-                    "unittest".to_owned(),
-                    VariableName::File.template_value(),
-                ],
-                ..TaskTemplate::default()
-            },
-            TaskTemplate {
-                label: "unittest $ZED_CUSTOM_PYTHON_UNITTEST_TARGET".to_owned(),
-                command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
-                args: vec![
-                    "-m".to_owned(),
-                    "unittest".to_owned(),
-                    "$ZED_CUSTOM_PYTHON_UNITTEST_TARGET".to_owned(),
-                ],
-                tags: vec![
-                    "python-unittest-class".to_owned(),
-                    "python-unittest-method".to_owned(),
-                ],
-                ..TaskTemplate::default()
-            },
-        ]))
+        ];
+
+        tasks.extend(match test_runner {
+            TestRunner::UNITTEST => {
+                [
+                    // Run tests for an entire file
+                    TaskTemplate {
+                        label: format!("unittest '{}'", VariableName::File.template_value()),
+                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
+                        args: vec![
+                            "-m".to_owned(),
+                            "unittest".to_owned(),
+                            VariableName::File.template_value(),
+                        ],
+                        ..TaskTemplate::default()
+                    },
+                    // Run test(s) for a specific target within a file
+                    TaskTemplate {
+                        label: "unittest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
+                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
+                        args: vec![
+                            "-m".to_owned(),
+                            "unittest".to_owned(),
+                            "$ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
+                        ],
+                        tags: vec![
+                            "python-unittest-class".to_owned(),
+                            "python-unittest-method".to_owned(),
+                        ],
+                        ..TaskTemplate::default()
+                    },
+                ]
+            }
+            TestRunner::PYTEST => {
+                [
+                    // Run tests for an entire file
+                    TaskTemplate {
+                        label: format!("pytest '{}'", VariableName::File.template_value()),
+                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
+                        args: vec![
+                            "-m".to_owned(),
+                            "pytest".to_owned(),
+                            VariableName::File.template_value(),
+                        ],
+                        ..TaskTemplate::default()
+                    },
+                    // Run test(s) for a specific target within a file
+                    TaskTemplate {
+                        label: "pytest $ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
+                        command: PYTHON_ACTIVE_TOOLCHAIN_PATH.template_value(),
+                        args: vec![
+                            "-m".to_owned(),
+                            "pytest".to_owned(),
+                            "$ZED_CUSTOM_PYTHON_TEST_TARGET".to_owned(),
+                        ],
+                        tags: vec![
+                            "python-pytest-class".to_owned(),
+                            "python-pytest-method".to_owned(),
+                        ],
+                        ..TaskTemplate::default()
+                    },
+                ]
+            }
+        });
+
+        Some(TaskTemplates(tasks))
+    }
+}
+
+fn selected_test_runner(location: Option<&Arc<dyn language::File>>, cx: &AppContext) -> TestRunner {
+    const TEST_RUNNER_VARIABLE: &str = "TEST_RUNNER";
+    language_settings(Some(LanguageName::new("Python")), location, cx)
+        .tasks
+        .variables
+        .get(TEST_RUNNER_VARIABLE)
+        .and_then(|val| TestRunner::from_str(val).ok())
+        .unwrap_or(TestRunner::PYTEST)
+}
+
+impl PythonContextProvider {
+    fn build_unittest_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 unittest_class_name =
+            variables.get(&VariableName::Custom(Cow::Borrowed("_unittest_class_name")));
+
+        let unittest_method_name = variables.get(&VariableName::Custom(Cow::Borrowed(
+            "_unittest_method_name",
+        )));
+
+        let unittest_target_str = match (unittest_class_name, unittest_method_name) {
+            (Some(class_name), Some(method_name)) => {
+                format!("{}.{}.{}", python_module_name, class_name, method_name)
+            }
+            (Some(class_name), None) => format!("{}.{}", python_module_name, class_name),
+            (None, None) => python_module_name,
+            (None, Some(_)) => return Ok((VariableName::Custom(Cow::Borrowed("")), String::new())), // should never happen, a TestCase class is the unit of testing
+        };
+
+        let unittest_target = (
+            PYTHON_TEST_TARGET_TASK_VARIABLE.clone(),
+            unittest_target_str,
+        );
+
+        Ok(unittest_target)
+    }
+
+    fn build_pytest_target(
+        &self,
+        variables: &task::TaskVariables,
+    ) -> Result<(VariableName, String)> {
+        let file_path = variables
+            .get(&VariableName::RelativeFile)
+            .ok_or_else(|| anyhow!("No file path given"))?;
+
+        let pytest_class_name =
+            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_class_name")));
+
+        let pytest_method_name =
+            variables.get(&VariableName::Custom(Cow::Borrowed("_pytest_method_name")));
+
+        let pytest_target_str = match (pytest_class_name, pytest_method_name) {
+            (Some(class_name), Some(method_name)) => {
+                format!("{}::{}::{}", file_path, class_name, method_name)
+            }
+            (Some(class_name), None) => {
+                format!("{}::{}", file_path, class_name)
+            }
+            (None, Some(method_name)) => {
+                format!("{}::{}", file_path, method_name)
+            }
+            (None, None) => file_path.to_string(),
+        };
+
+        let pytest_target = (PYTHON_TEST_TARGET_TASK_VARIABLE.clone(), pytest_target_str);
+
+        Ok(pytest_target)
     }
 }
 

crates/languages/src/python/runnables.scm

@@ -29,3 +29,42 @@
             )
         )
 )
+
+; pytest functions
+(
+    (module
+        (function_definition
+            name: (identifier) @run @_pytest_method_name
+            (#match? @_pytest_method_name "^test_")
+        ) @python-pytest-method
+    )
+    (#set! tag python-pytest-method)
+)
+
+; pytest classes
+(
+    (module
+        (class_definition
+            name: (identifier) @run @_pytest_class_name
+            (#match? @_pytest_class_name "^Test")
+        )
+        (#set! tag python-pytest-class)
+    )
+)
+
+; pytest class methods
+(
+    (module
+        (class_definition
+            name: (identifier) @_pytest_class_name
+            (#match? @_pytest_class_name "^Test")
+            body: (block
+                    (function_definition
+                        name: (identifier) @run @_pytest_method_name
+                        (#match? @_pytest_method_name "^test")
+                    ) @python-pytest-method
+                    (#set! tag python-pytest-method)
+            )
+        )
+    )
+)