1use std::path::Path;
  2
  3use anyhow::{Result, bail};
  4use async_trait::async_trait;
  5use dap::{DapLocator, DebugRequest, adapters::DebugAdapterName};
  6use gpui::SharedString;
  7
  8use task::{DebugScenario, SpawnInTerminal, TaskTemplate, VariableName};
  9
 10pub(crate) struct PythonLocator;
 11
 12#[async_trait]
 13impl DapLocator for PythonLocator {
 14    fn name(&self) -> SharedString {
 15        SharedString::new_static("Python")
 16    }
 17
 18    /// Determines whether this locator can generate debug target for given task.
 19    async fn create_scenario(
 20        &self,
 21        build_config: &TaskTemplate,
 22        resolved_label: &str,
 23        adapter: &DebugAdapterName,
 24    ) -> Option<DebugScenario> {
 25        if adapter.0.as_ref() != "Debugpy" {
 26            return None;
 27        }
 28        let valid_program = build_config.command.starts_with("$ZED_")
 29            || Path::new(&build_config.command)
 30                .file_name()
 31                .is_some_and(|name| name.to_str().is_some_and(|path| path.starts_with("python")));
 32        if !valid_program || build_config.args.iter().any(|arg| arg == "-c") {
 33            // We cannot debug selections.
 34            return None;
 35        }
 36        let command = build_config.command.clone();
 37        let module_specifier_position = build_config
 38            .args
 39            .iter()
 40            .position(|arg| arg == "-m")
 41            .map(|position| position + 1);
 42        // Skip the -m and module name, get all that's after.
 43        let mut rest_of_the_args = module_specifier_position
 44            .and_then(|position| build_config.args.get(position..))
 45            .into_iter()
 46            .flatten()
 47            .fuse();
 48        let mod_name = rest_of_the_args.next();
 49        let args = rest_of_the_args.collect::<Vec<_>>();
 50
 51        let program_position = mod_name
 52            .is_none()
 53            .then(|| {
 54                let zed_file = VariableName::File.template_value_with_whitespace();
 55                build_config.args.iter().position(|arg| *arg == zed_file)
 56            })
 57            .flatten();
 58        let args = if let Some(position) = program_position {
 59            args.into_iter().skip(position).collect::<Vec<_>>()
 60        } else {
 61            args
 62        };
 63        if program_position.is_none() && mod_name.is_none() {
 64            return None;
 65        }
 66        let mut config = serde_json::json!({
 67            "request": "launch",
 68            "python": command,
 69            "args": args,
 70            "cwd": build_config.cwd.clone()
 71        });
 72        if let Some(config_obj) = config.as_object_mut() {
 73            if let Some(module) = mod_name {
 74                config_obj.insert("module".to_string(), module.clone().into());
 75            }
 76            if let Some(program) = program_position {
 77                config_obj.insert(
 78                    "program".to_string(),
 79                    build_config.args[program].clone().into(),
 80                );
 81            }
 82        }
 83
 84        Some(DebugScenario {
 85            adapter: adapter.0.clone(),
 86            label: resolved_label.to_string().into(),
 87            build: None,
 88            config,
 89            tcp_connection: None,
 90        })
 91    }
 92
 93    async fn run(&self, _: SpawnInTerminal) -> Result<DebugRequest> {
 94        bail!("Python locator should not require DapLocator::run to be ran");
 95    }
 96}
 97
 98#[cfg(test)]
 99mod test {
100    use serde_json::json;
101
102    use super::*;
103
104    #[gpui::test]
105    async fn test_python_locator() {
106        let adapter = DebugAdapterName("Debugpy".into());
107        let build_task = TaskTemplate {
108            label: "run module '$ZED_FILE'".into(),
109            command: "$ZED_CUSTOM_PYTHON_ACTIVE_ZED_TOOLCHAIN".into(),
110            args: vec!["-m".into(), "$ZED_CUSTOM_PYTHON_MODULE_NAME".into()],
111            env: Default::default(),
112            cwd: Some("$ZED_WORKTREE_ROOT".into()),
113            use_new_terminal: false,
114            allow_concurrent_runs: false,
115            reveal: task::RevealStrategy::Always,
116            reveal_target: task::RevealTarget::Dock,
117            hide: task::HideStrategy::Never,
118            tags: vec!["python-module-main-method".into()],
119            languages: vec![],
120            shell: task::Shell::System,
121            show_summary: false,
122            show_command: false,
123        };
124
125        let expected_scenario = DebugScenario {
126            adapter: "Debugpy".into(),
127            label: "run module 'main.py'".into(),
128            build: None,
129            config: json!({
130                "request": "launch",
131                "python": "$ZED_CUSTOM_PYTHON_ACTIVE_ZED_TOOLCHAIN",
132                "args": [],
133                "cwd": "$ZED_WORKTREE_ROOT",
134                "module": "$ZED_CUSTOM_PYTHON_MODULE_NAME",
135            }),
136            tcp_connection: None,
137        };
138
139        assert_eq!(
140            PythonLocator
141                .create_scenario(&build_task, "run module 'main.py'", &adapter)
142                .await
143                .expect("Failed to create a scenario"),
144            expected_scenario
145        );
146    }
147}