Skip to main content

usage/spec/
choices.rs

1#[cfg(feature = "unstable_choices_env")]
2use kdl::KdlEntry;
3use kdl::KdlNode;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7use crate::error::UsageErr;
8use crate::spec::context::ParsingContext;
9use crate::spec::helpers::NodeHelper;
10
11#[derive(Debug, Default, Clone, Serialize, Deserialize)]
12#[non_exhaustive]
13pub struct SpecChoices {
14    pub choices: Vec<String>,
15    #[cfg(feature = "unstable_choices_env")]
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub env: Option<String>,
18}
19
20impl SpecChoices {
21    /// The set of values an arg or flag accepts.
22    pub fn new(choices: impl IntoIterator<Item = impl Into<String>>) -> Self {
23        Self {
24            choices: choices.into_iter().map(Into::into).collect(),
25            ..Default::default()
26        }
27    }
28}
29
30impl SpecChoices {
31    #[cfg(feature = "unstable_choices_env")]
32    #[must_use]
33    pub fn env(&self) -> Option<&str> {
34        self.env.as_deref()
35    }
36
37    #[cfg(not(feature = "unstable_choices_env"))]
38    #[must_use]
39    pub fn env(&self) -> Option<&str> {
40        None
41    }
42
43    #[cfg(feature = "unstable_choices_env")]
44    pub fn set_env(&mut self, env: Option<String>) {
45        self.env = env;
46    }
47
48    pub(crate) fn parse(ctx: &ParsingContext, node: &NodeHelper) -> Result<Self, UsageErr> {
49        #[cfg(not(feature = "unstable_choices_env"))]
50        node.ensure_arg_len(1..)?;
51
52        #[cfg(feature = "unstable_choices_env")]
53        let mut config = Self {
54            choices: node
55                .args()
56                .map(|e| e.ensure_string())
57                .collect::<Result<_, _>>()?,
58            ..Default::default()
59        };
60
61        #[cfg(not(feature = "unstable_choices_env"))]
62        let config = Self {
63            choices: node
64                .args()
65                .map(|e| e.ensure_string())
66                .collect::<Result<_, _>>()?,
67            ..Default::default()
68        };
69
70        for (k, v) in node.props() {
71            match k {
72                #[cfg(feature = "unstable_choices_env")]
73                "env" => config.set_env(Some(v.ensure_string()?)),
74                k => bail_parse!(ctx, v.entry.span(), "unsupported choices key {k}"),
75            }
76        }
77
78        if config.choices.is_empty() {
79            #[cfg(feature = "unstable_choices_env")]
80            if config.env().is_none() {
81                bail_parse!(
82                    ctx,
83                    node.span(),
84                    "choices must have at least 1 argument or env property"
85                );
86            }
87            #[cfg(not(feature = "unstable_choices_env"))]
88            bail_parse!(ctx, node.span(), "choices must have at least 1 argument");
89        }
90
91        Ok(config)
92    }
93
94    pub fn values(&self) -> Vec<String> {
95        self.values_with_env(None)
96    }
97
98    pub(crate) fn values_with_env(&self, env: Option<&HashMap<String, String>>) -> Vec<String> {
99        #[cfg(feature = "unstable_choices_env")]
100        let mut values = self.choices.clone();
101
102        #[cfg(not(feature = "unstable_choices_env"))]
103        let values = self.choices.clone();
104
105        #[cfg(not(feature = "unstable_choices_env"))]
106        let _ = env;
107
108        #[cfg(feature = "unstable_choices_env")]
109        {
110            if let Some(env_key) = self.env() {
111                let env_value = if let Some(env_map) = env {
112                    env_map.get(env_key).cloned()
113                } else {
114                    std::env::var(env_key).ok()
115                };
116
117                if let Some(env_value) = env_value {
118                    for choice in env_value
119                        .split(|c: char| c == ',' || c.is_whitespace())
120                        .filter(|choice| !choice.is_empty())
121                    {
122                        let choice = choice.to_string();
123                        if !values.contains(&choice) {
124                            values.push(choice);
125                        }
126                    }
127                }
128            }
129        }
130
131        values
132    }
133}
134
135impl From<&SpecChoices> for KdlNode {
136    fn from(arg: &SpecChoices) -> Self {
137        let mut node = KdlNode::new("choices");
138        for choice in &arg.choices {
139            node.push(choice.to_string());
140        }
141        #[cfg(feature = "unstable_choices_env")]
142        if let Some(env) = arg.env() {
143            node.push(KdlEntry::new_prop("env", env.to_string()));
144        }
145        node
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    #[cfg(feature = "unstable_choices_env")]
152    use super::SpecChoices;
153    #[cfg(feature = "unstable_choices_env")]
154    use std::collections::HashMap;
155
156    #[cfg(feature = "unstable_choices_env")]
157    #[test]
158    fn values_with_env_splits_on_commas_and_whitespace() {
159        let mut choices = SpecChoices {
160            choices: vec!["local".into()],
161            ..Default::default()
162        };
163        choices.set_env(Some("DEPLOY_ENVS".into()));
164
165        let env = HashMap::from([("DEPLOY_ENVS".to_string(), "foo,bar baz\nqux".to_string())]);
166
167        assert_eq!(
168            choices.values_with_env(Some(&env)),
169            vec!["local", "foo", "bar", "baz", "qux"]
170        );
171    }
172
173    #[cfg(feature = "unstable_choices_env")]
174    #[test]
175    fn values_with_env_deduplicates_existing_choices() {
176        let mut choices = SpecChoices {
177            choices: vec!["foo".into()],
178            ..Default::default()
179        };
180        choices.set_env(Some("DEPLOY_ENVS".into()));
181
182        let env = HashMap::from([("DEPLOY_ENVS".to_string(), "foo,bar foo".to_string())]);
183
184        assert_eq!(choices.values_with_env(Some(&env)), vec!["foo", "bar"]);
185    }
186
187    #[cfg(feature = "unstable_choices_env")]
188    #[test]
189    fn values_with_env_does_not_fallback_when_custom_env_is_present() {
190        let mut choices = SpecChoices {
191            choices: vec!["local".into()],
192            ..Default::default()
193        };
194        choices.set_env(Some(
195            "USAGE_TEST_CHOICES_ENV_DOES_NOT_EXIST_A5E0F4D1".into(),
196        ));
197
198        assert_eq!(
199            choices.values_with_env(Some(&HashMap::new())),
200            vec!["local"]
201        );
202    }
203}