feat: disambiguate tasks interactive

Cargo.toml

@@ -26,6 +26,7 @@ clap-verbosity-flag = "2.1.2"
 clap_complete = "4.4.9"
 console = { version = "0.15.8", features = ["windows-console-colors"] }
 deno_task_shell = "0.14.3"
+dialoguer = "0.11.0"
 dirs = "5.0.1"
 dunce = "1.0.4"
 flate2 = "1.0.28"

src/cli/run.rs

@@ -1,16 +1,21 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashSet;
+use std::convert::identity;
 use std::str::FromStr;
 use std::{collections::HashMap, path::PathBuf, string::String};
 
 use clap::Parser;
+use dialoguer::theme::ColorfulTheme;
 use itertools::Itertools;
 use miette::{miette, Context, Diagnostic};
 use rattler_conda_types::Platform;
 
 use crate::activation::get_environment_variables;
 use crate::project::errors::UnsupportedPlatformError;
-use crate::task::{ExecutableTask, FailedToParseShellScript, InvalidWorkingDirectory, TaskGraph};
+use crate::task::{
+    AmbiguousTask, ExecutableTask, FailedToParseShellScript, InvalidWorkingDirectory,
+    SearchEnvironments, TaskAndEnvironment, TaskGraph,
+};
 use crate::{Project, UpdateLockFileOptions};
 
 use crate::environment::LockFileDerivedData;
@@ -77,12 +82,14 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     tracing::debug!("Task parsed from run command: {:?}", task_args);
 
     // Construct a task graph from the input arguments
-    let task_graph = TaskGraph::from_cmd_args(
+    let search_environment = SearchEnvironments::from_opt_env(
         &project,
-        task_args,
-        Some(Platform::current()),
         explicit_environment.clone(),
-    )?;
+        Some(Platform::current()),
+    )
+    .with_disambiguate_fn(disambiguate_task_interactive);
+
+    let task_graph = TaskGraph::from_cmd_args(&project, &search_environment, task_args)?;
 
     // Traverse the task graph in topological order and execute each individual task.
     let mut task_idx = 0;
@@ -254,3 +261,30 @@ async fn execute_task<'p>(
 
     Ok(())
 }
+
+/// Called to disambiguate between environments to run a task in.
+fn disambiguate_task_interactive<'p>(
+    problem: &AmbiguousTask<'p>,
+) -> Option<TaskAndEnvironment<'p>> {
+    let environment_names = problem
+        .environments
+        .iter()
+        .map(|(env, _)| env.name())
+        .collect_vec();
+    dialoguer::Select::with_theme(&ColorfulTheme::default())
+        .with_prompt(format!(
+            "The task '{}' {}can be run in multiple environments.\n\nPlease select an environment to run the task in:",
+            problem.task_name,
+            if let Some(dependency) = &problem.depended_on_by {
+                format!("(depended on by '{}') ", dependency.0)
+            } else {
+                String::new()
+            }
+        ))
+        .report(false)
+        .items(&environment_names)
+        .default(0)
+        .interact_opt()
+        .map_or(None, identity)
+        .map(|idx| problem.environments[idx].clone())
+}

src/lib.rs

@@ -28,7 +28,8 @@ pub use project::{
     DependencyType, Project, SpecType,
 };
 pub use task::{
-    CmdArgs, ExecutableTask, RunOutput, Task, TaskExecutionError, TaskGraph, TaskGraphError,
+    CmdArgs, ExecutableTask, FindTaskError, FindTaskSource, RunOutput, SearchEnvironments, Task,
+    TaskDisambiguation, TaskExecutionError, TaskGraph, TaskGraphError,
 };
 
 use rattler_networking::retry_policies::ExponentialBackoff;

src/task/mod.rs

@@ -7,12 +7,17 @@ use std::path::{Path, PathBuf};
 
 mod error;
 mod executable_task;
+mod task_environment;
 mod task_graph;
 
 pub use executable_task::{
     ExecutableTask, FailedToParseShellScript, InvalidWorkingDirectory, RunOutput,
     TaskExecutionError,
 };
+pub use task_environment::{
+    AmbiguousTask, FindTaskError, FindTaskSource, SearchEnvironments, TaskAndEnvironment,
+    TaskDisambiguation,
+};
 pub use task_graph::{TaskGraph, TaskGraphError, TaskId, TaskNode};
 
 /// Represents different types of scripts
@@ -22,7 +27,8 @@ pub enum Task {
     Plain(String),
     Execute(Execute),
     Alias(Alias),
-    // We don't what a way for the deserializer to except a custom task, as they are meant for tasks given in the command line.
+    // We want a way for the deserializer to except a custom task, as they are meant for tasks
+    // given in the command line.
     #[serde(skip)]
     Custom(Custom),
 }

src/task/task_environment.rs

@@ -0,0 +1,191 @@
+use crate::project::Environment;
+use crate::task::error::{AmbiguousTaskError, MissingTaskError};
+use crate::{Project, Task};
+use itertools::Itertools;
+use miette::Diagnostic;
+use rattler_conda_types::Platform;
+use thiserror::Error;
+
+/// Defines where the task was defined when looking for a task.
+#[derive(Debug, Clone)]
+pub enum FindTaskSource<'p> {
+    CmdArgs,
+    DependsOn(String, &'p Task),
+}
+
+pub type TaskAndEnvironment<'p> = (Environment<'p>, &'p Task);
+
+pub trait TaskDisambiguation<'p> {
+    fn disambiguate(&self, task: &AmbiguousTask<'p>) -> Option<TaskAndEnvironment<'p>>;
+}
+
+#[derive(Default)]
+pub struct NoDisambiguation;
+pub struct DisambiguateFn<Fn>(Fn);
+
+impl<'p> TaskDisambiguation<'p> for NoDisambiguation {
+    fn disambiguate(&self, _task: &AmbiguousTask<'p>) -> Option<TaskAndEnvironment<'p>> {
+        None
+    }
+}
+
+impl<'p, F: Fn(&AmbiguousTask<'p>) -> Option<TaskAndEnvironment<'p>>> TaskDisambiguation<'p>
+    for DisambiguateFn<F>
+{
+    fn disambiguate(&self, task: &AmbiguousTask<'p>) -> Option<TaskAndEnvironment<'p>> {
+        self.0(task)
+    }
+}
+
+/// An object to help with searching for tasks.
+pub struct SearchEnvironments<'p, D: TaskDisambiguation<'p> = NoDisambiguation> {
+    pub project: &'p Project,
+    pub explicit_environment: Option<Environment<'p>>,
+    pub platform: Option<Platform>,
+    pub disambiguate: D,
+}
+
+/// Information about an task that was found when searching for a task
+pub struct AmbiguousTask<'p> {
+    pub task_name: String,
+    pub depended_on_by: Option<(String, &'p Task)>,
+    pub environments: Vec<TaskAndEnvironment<'p>>,
+}
+
+impl<'p> From<AmbiguousTask<'p>> for AmbiguousTaskError {
+    fn from(value: AmbiguousTask<'p>) -> Self {
+        Self {
+            task_name: value.task_name,
+            environments: value
+                .environments
+                .into_iter()
+                .map(|env| env.0.name().clone())
+                .collect(),
+        }
+    }
+}
+
+#[derive(Debug, Diagnostic, Error)]
+pub enum FindTaskError {
+    #[error(transparent)]
+    MissingTask(MissingTaskError),
+
+    #[error(transparent)]
+    AmbiguousTask(AmbiguousTaskError),
+}
+
+impl<'p> SearchEnvironments<'p, NoDisambiguation> {
+    // Determine which environments we are allowed to check for tasks.
+    //
+    // If the user specified an environment, look for tasks in the main environment and the
+    // user specified environment.
+    //
+    // If the user did not specify an environment, look for tasks in any environment.
+    pub fn from_opt_env(
+        project: &'p Project,
+        explicit_environment: Option<Environment<'p>>,
+        platform: Option<Platform>,
+    ) -> Self {
+        Self {
+            project,
+            explicit_environment,
+            platform,
+            disambiguate: NoDisambiguation,
+        }
+    }
+}
+
+impl<'p, D: TaskDisambiguation<'p>> SearchEnvironments<'p, D> {
+    /// Returns a new `SearchEnvironments` with the given disambiguation function.
+    pub fn with_disambiguate_fn<F: Fn(&AmbiguousTask<'p>) -> Option<TaskAndEnvironment<'p>>>(
+        self,
+        func: F,
+    ) -> SearchEnvironments<'p, DisambiguateFn<F>> {
+        SearchEnvironments {
+            project: self.project,
+            explicit_environment: self.explicit_environment,
+            platform: self.platform,
+            disambiguate: DisambiguateFn(func),
+        }
+    }
+
+    /// Finds the task with the given name or returns an error that explains why the task could not
+    /// be found.
+    pub fn find_task(
+        &self,
+        name: &str,
+        source: FindTaskSource<'p>,
+    ) -> Result<TaskAndEnvironment<'p>, FindTaskError> {
+        // If the task was specified on the command line and there is no explicit environment and
+        // the task is only defined in the default feature, use the default environment.
+        if matches!(source, FindTaskSource::CmdArgs) && self.explicit_environment.is_none() {
+            if let Some(task) = self
+                .project
+                .manifest
+                .default_feature()
+                .targets
+                .resolve(self.platform)
+                .find_map(|target| target.tasks.get(name))
+            {
+                // None of the other environments can have this task. Otherwise, its still
+                // ambiguous.
+                if !self
+                    .project
+                    .environments()
+                    .into_iter()
+                    .flat_map(|env| env.features(false).collect_vec())
+                    .flat_map(|feature| feature.targets.resolve(self.platform))
+                    .any(|target| target.tasks.contains_key(name))
+                {
+                    return Ok((self.project.default_environment(), task));
+                }
+            }
+        }
+
+        // If an explicit environment was specified, only look for tasks in that environment and
+        // the default environment.
+        let environments = if let Some(explicit_environment) = &self.explicit_environment {
+            vec![explicit_environment.clone()]
+        } else {
+            self.project.environments()
+        };
+
+        // Find all the task and environment combinations
+        let include_default_feature = true;
+        let mut tasks = Vec::new();
+        for env in environments.iter() {
+            if let Some(task) = env
+                .tasks(self.platform, include_default_feature)
+                .ok()
+                .and_then(|tasks| tasks.get(name).copied())
+            {
+                tasks.push((env.clone(), task));
+            }
+        }
+
+        match tasks.len() {
+            0 => Err(FindTaskError::MissingTask(MissingTaskError {
+                task_name: name.to_string(),
+            })),
+            1 => {
+                let (env, task) = tasks.remove(0);
+                Ok((env.clone(), task))
+            }
+            _ => {
+                let ambiguous_task = AmbiguousTask {
+                    task_name: name.to_string(),
+                    depended_on_by: match source {
+                        FindTaskSource::DependsOn(dep, task) => Some((dep, task)),
+                        _ => None,
+                    },
+                    environments: tasks,
+                };
+
+                match self.disambiguate.disambiguate(&ambiguous_task) {
+                    Some(env) => Ok(env),
+                    None => Err(FindTaskError::AmbiguousTask(ambiguous_task.into())),
+                }
+            }
+        }
+    }
+}