Skip to main content

shared/extensions/
commands.rs

1use clap::{Arg, ArgMatches, Args, Command};
2use std::{collections::HashMap, pin::Pin, sync::Arc};
3
4pub type ExecutorFunc = dyn Fn(
5        Option<Arc<crate::env::Env>>,
6        ArgMatches,
7    ) -> Pin<Box<dyn Future<Output = Result<i32, anyhow::Error>>>>
8    + Send;
9
10pub enum CommandMapEntry {
11    Command(Box<ExecutorFunc>),
12    Group(HashMap<&'static str, CommandMapEntry>),
13}
14
15pub struct CliCommandGroupBuilder {
16    command: Command,
17    map: HashMap<&'static str, CommandMapEntry>,
18}
19
20impl CliCommandGroupBuilder {
21    pub fn new(name: &'static str, about: &'static str) -> Self {
22        Self {
23            command: Command::new(name)
24                .version(crate::full_version())
25                .arg(
26                    Arg::new("debug")
27                        .help("pass in order to run in debug mode")
28                        .num_args(0)
29                        .short('d')
30                        .long("debug")
31                        .default_value("false")
32                        .value_parser(clap::value_parser!(bool))
33                        .global(true)
34                        .required(false),
35                )
36                .about(about),
37            map: HashMap::new(),
38        }
39    }
40
41    /// Moves every command and group of `other` into this group, unless one of them is already registered.
42    pub(crate) fn merge(&mut self, other: CliCommandGroupBuilder) -> Result<(), String> {
43        if let Some(name) = other.map.keys().find(|name| self.map.contains_key(*name)) {
44            return Err(format!("CLI command `{name}` is already registered"));
45        }
46
47        let mut command = std::mem::take(&mut self.command);
48        for subcommand in other.command.get_subcommands() {
49            command = command.subcommand(subcommand.clone());
50        }
51        self.command = command;
52        self.map.extend(other.map);
53
54        Ok(())
55    }
56
57    pub fn get_command(&self) -> Command {
58        self.command.clone()
59    }
60
61    pub fn get_matches(&mut self) -> ArgMatches {
62        self.command.get_matches_mut()
63    }
64
65    pub fn print_help(&mut self) {
66        let _ = self.command.print_long_help();
67    }
68
69    pub fn match_command(
70        &self,
71        command: String,
72        arg_matches: ArgMatches,
73    ) -> Option<(&ExecutorFunc, ArgMatches)> {
74        let mut current_map = &self.map;
75        let mut current_matches = arg_matches;
76        let mut current_command = command;
77
78        loop {
79            let entry = current_map.get(current_command.as_str())?;
80
81            match entry {
82                CommandMapEntry::Command(executor) => {
83                    return Some((executor, current_matches));
84                }
85                CommandMapEntry::Group(submap) => {
86                    let (subcommand_name, subcommand_matches) =
87                        current_matches.remove_subcommand()?;
88
89                    current_map = submap;
90                    current_matches = subcommand_matches;
91                    current_command = subcommand_name;
92                }
93            }
94        }
95    }
96
97    pub fn add_group<F: FnOnce(CliCommandGroupBuilder) -> CliCommandGroupBuilder>(
98        mut self,
99        name: &'static str,
100        about: &'static str,
101        callback: F,
102    ) -> Self {
103        let subgroup = CliCommandGroupBuilder::new(name, about);
104        let subgroup = callback(subgroup);
105
106        self.command = self.command.subcommand(subgroup.command);
107        self.map.insert(name, CommandMapEntry::Group(subgroup.map));
108
109        self
110    }
111
112    pub fn add_command<A: Args>(
113        mut self,
114        name: &'static str,
115        about: &'static str,
116        cli_command: impl CliCommand<A>,
117    ) -> Self {
118        let command = cli_command.get_command(Command::new(name).about(about));
119        let command = A::augment_args(command);
120
121        self.command = self.command.subcommand(command);
122        self.map
123            .insert(name, CommandMapEntry::Command(cli_command.get_executor()));
124
125        self
126    }
127}
128
129pub trait CliCommand<A: Args> {
130    fn get_command(&self, command: Command) -> Command;
131    fn get_executor(self) -> Box<ExecutorFunc>;
132}