Skip to main content

bake/
arguments.rs

1use crate::{Error, Result};
2use std::collections::BTreeMap;
3use std::fmt::Display;
4use std::str::FromStr;
5
6/// A task parameter. Names use underscores; command-line flags also accept hyphens.
7pub struct Parameter {
8    pub(crate) name: String,
9    pub(crate) type_name: &'static str,
10    pub(crate) positional: bool,
11    pub(crate) required: bool,
12    pub(crate) repeated: bool,
13    pub(crate) default: Option<String>,
14    pub(crate) description: String,
15    validate: fn(&str) -> Result<()>,
16}
17
18impl Parameter {
19    /// Values use `FromStr`; custom types can implement it too.
20    pub fn new<Value: FromStr>(name: &str) -> Self
21    where
22        Value::Err: Display,
23    {
24        Self {
25            name: name.into(),
26            type_name: std::any::type_name::<Value>(),
27            positional: true,
28            required: true,
29            repeated: false,
30            default: None,
31            description: String::new(),
32            validate: |value| parse::<Value>(value).map(|_| ()),
33        }
34    }
35    pub fn named(mut self) -> Self {
36        self.positional = false;
37        self
38    }
39    pub fn optional(mut self) -> Self {
40        self.required = false;
41        self
42    }
43    pub fn repeated(mut self) -> Self {
44        self.repeated = true;
45        self.positional = false;
46        self.required = false;
47        self
48    }
49    pub fn default(mut self, value: impl Into<String>) -> Self {
50        self.default = Some(value.into());
51        self.required = false;
52        self.positional = false;
53        self
54    }
55    pub fn help(mut self, description: impl Into<String>) -> Self {
56        self.description = description.into();
57        self
58    }
59    pub fn name(&self) -> &str {
60        &self.name
61    }
62}
63
64fn parse<Value: FromStr>(value: &str) -> Result<Value>
65where
66    Value::Err: Display,
67{
68    value.parse().map_err(|error| {
69        Error::new(format!(
70            "invalid {} value {value:?}: {error}",
71            std::any::type_name::<Value>()
72        ))
73    })
74}
75
76/// Validated command-line values, converted to Rust types by the task adapter.
77#[derive(Default, Debug)]
78pub struct Arguments(BTreeMap<String, Vec<String>>);
79
80impl Arguments {
81    pub fn required<Value: FromStr>(&self, name: &str) -> Result<Value>
82    where
83        Value::Err: Display,
84    {
85        self.optional(name)?
86            .ok_or_else(|| Error::new(format!("missing argument {name:?}")))
87    }
88    pub fn optional<Value: FromStr>(&self, name: &str) -> Result<Option<Value>>
89    where
90        Value::Err: Display,
91    {
92        self.0
93            .get(name)
94            .and_then(|values| values.first())
95            .map(|value| parse(value))
96            .transpose()
97    }
98    pub fn repeated<Value: FromStr>(&self, name: &str) -> Result<Vec<Value>>
99    where
100        Value::Err: Display,
101    {
102        self.0
103            .get(name)
104            .into_iter()
105            .flatten()
106            .map(|value| parse(value))
107            .collect()
108    }
109
110    fn insert(&mut self, parameter: &Parameter, value: &str) -> Result<()> {
111        if !parameter.repeated && self.0.contains_key(&parameter.name) {
112            return Err(Error::new(format!(
113                "argument {:?} was supplied more than once",
114                parameter.name
115            )));
116        }
117        (parameter.validate)(value)
118            .map_err(|error| Error::new(format!("{}: {error}", parameter.name)))?;
119        self.0
120            .entry(parameter.name.clone())
121            .or_default()
122            .push(value.into());
123        Ok(())
124    }
125
126    pub(crate) fn extract(parameters: &[Parameter], tokens: &[String]) -> Result<(Self, usize)> {
127        let mut arguments = Self::default();
128        let mut consumed = 0;
129        let mut options = true;
130        while let Some(token) = tokens.get(consumed) {
131            if token == "::" {
132                break;
133            }
134            if options && token == "--" {
135                options = false;
136                consumed += 1;
137                continue;
138            }
139            let named = if options {
140                token.strip_prefix("--")
141            } else {
142                None
143            };
144            if let Some(name) = named {
145                if name.contains('=') {
146                    return Err(Error::new(
147                        "named arguments use separate `--name value` tokens",
148                    ));
149                }
150                let name = name.replace('-', "_");
151                let parameter = parameters
152                    .iter()
153                    .find(|parameter| parameter.name == name)
154                    .ok_or_else(|| Error::new(format!("unknown argument {name:?}")))?;
155                consumed += 1;
156                let value = tokens
157                    .get(consumed)
158                    .filter(|value| value.as_str() != "::" && value.as_str() != "--")
159                    .ok_or_else(|| Error::new(format!("argument {name:?} requires a value")))?;
160                consumed += 1;
161                arguments.insert(parameter, value)?;
162            } else if let Some(parameter) = parameters.iter().find(|parameter| {
163                parameter.positional && !arguments.0.contains_key(&parameter.name)
164            }) {
165                arguments.insert(parameter, token)?;
166                consumed += 1;
167            } else {
168                break;
169            }
170        }
171        for parameter in parameters {
172            if parameter.required && !arguments.0.contains_key(&parameter.name) {
173                return Err(Error::new(format!("missing argument {:?}", parameter.name)));
174            }
175        }
176        Ok((arguments, consumed))
177    }
178}