Skip to main content

command_stream/zx/
argv.rs

1//! Command line parsing compatible with `minimist` (as re-exported by zx) and
2//! zx's `parseArgv` wrapper.
3//!
4//! Results are returned as a [`serde_json::Value`] object so that repeated
5//! flags (arrays), dotted keys (nested objects), numbers, booleans and strings
6//! can all be represented. Positional arguments live under `"_"`.
7
8use std::collections::{HashMap, HashSet};
9
10use once_cell::sync::Lazy;
11use regex::Regex;
12use serde_json::{Map, Value};
13
14use super::util::to_camel_case;
15
16static NUMBER_RE: Lazy<Regex> =
17    Lazy::new(|| Regex::new(r"^[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[-+]?\d+)?$").expect("valid regex"));
18static HEX_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^(?i)0x[0-9a-f]+$").expect("valid regex"));
19static FLAG_LIKE_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^(-|--)[^-]").expect("valid regex"));
20static NUMERIC_TAIL_RE: Lazy<Regex> =
21    Lazy::new(|| Regex::new(r"-?\d+(\.\d*)?(e-?\d+)?$").expect("valid regex"));
22const UNSAFE_KEYS: [&str; 3] = ["__proto__", "constructor", "prototype"];
23
24/// Which flags are booleans.
25#[derive(Debug, Clone, Default, PartialEq, Eq)]
26pub enum Booleans {
27    /// No flag is forced to be boolean.
28    #[default]
29    None,
30    /// Every `--flag` without `=` is boolean (`boolean: true`).
31    All,
32    /// The listed flags are boolean.
33    Keys(Vec<String>),
34}
35
36/// Callback for flags and positionals that were not declared; returning
37/// `false` drops the argument.
38pub type UnknownFn = fn(&str) -> bool;
39
40/// Options for [`minimist`] and [`parse_argv`].
41#[derive(Debug, Clone, Default)]
42pub struct ArgvOptions {
43    /// Boolean flags.
44    pub boolean: Booleans,
45    /// Flags whose values are always kept as strings (`"_"` for positionals).
46    pub string: Vec<String>,
47    /// Alias groups: key -> other names.
48    pub alias: Vec<(String, Vec<String>)>,
49    /// Default values (dotted keys allowed).
50    pub default: Vec<(String, Value)>,
51    /// Stop parsing at the first positional argument.
52    pub stop_early: bool,
53    /// Put arguments after `--` into `"--"` instead of `"_"`.
54    pub double_dash: bool,
55    /// Called for undeclared arguments.
56    pub unknown: Option<UnknownFn>,
57    /// (`parse_argv` only) convert `kebab-case` keys to `camelCase`.
58    pub camel_case: bool,
59    /// (`parse_argv` only) turn `"true"`/`"false"` string values into booleans.
60    pub parse_boolean: bool,
61}
62
63impl ArgvOptions {
64    /// Empty options.
65    pub fn new() -> Self {
66        Self::default()
67    }
68
69    /// Declare boolean flags.
70    pub fn boolean<S: AsRef<str>>(mut self, keys: &[S]) -> Self {
71        self.boolean = Booleans::Keys(keys.iter().map(|k| k.as_ref().to_string()).collect());
72        self
73    }
74
75    /// Treat every long flag without `=` as boolean.
76    pub fn all_boolean(mut self) -> Self {
77        self.boolean = Booleans::All;
78        self
79    }
80
81    /// Declare string flags.
82    pub fn string<S: AsRef<str>>(mut self, keys: &[S]) -> Self {
83        self.string = keys.iter().map(|k| k.as_ref().to_string()).collect();
84        self
85    }
86
87    /// Add an alias group.
88    pub fn alias<S: AsRef<str>>(mut self, key: &str, names: &[S]) -> Self {
89        let names = names.iter().map(|n| n.as_ref().to_string()).collect();
90        self.alias.push((key.to_string(), names));
91        self
92    }
93
94    /// Add a default value.
95    pub fn default_value(mut self, key: &str, value: impl Into<Value>) -> Self {
96        self.default.push((key.to_string(), value.into()));
97        self
98    }
99
100    /// Enable `stop_early`.
101    pub fn stop_early(mut self) -> Self {
102        self.stop_early = true;
103        self
104    }
105
106    /// Enable the `"--"` key.
107    pub fn double_dash(mut self) -> Self {
108        self.double_dash = true;
109        self
110    }
111
112    /// Set the unknown-argument callback.
113    pub fn unknown(mut self, callback: UnknownFn) -> Self {
114        self.unknown = Some(callback);
115        self
116    }
117
118    /// Enable camelCase keys (`parse_argv`).
119    pub fn camel_case(mut self) -> Self {
120        self.camel_case = true;
121        self
122    }
123
124    /// Enable boolean parsing of values (`parse_argv`).
125    pub fn parse_boolean(mut self) -> Self {
126        self.parse_boolean = true;
127        self
128    }
129}
130
131fn is_number(text: &str) -> bool {
132    HEX_RE.is_match(text) || NUMBER_RE.is_match(text)
133}
134
135/// Convert numeric text the way JavaScript's `Number()` does for the
136/// accepted syntaxes; integral values become integers.
137fn to_number(text: &str) -> Value {
138    let value = if HEX_RE.is_match(text) {
139        u64::from_str_radix(&text[2..], 16).map_or(f64::NAN, |v| v as f64)
140    } else {
141        text.parse::<f64>().unwrap_or(f64::NAN)
142    };
143    number_value(value)
144}
145
146fn number_value(value: f64) -> Value {
147    if value.fract() == 0.0 && value.abs() < 9.0e15 {
148        Value::from(value as i64)
149    } else {
150        serde_json::Number::from_f64(value).map_or(Value::Null, Value::Number)
151    }
152}
153
154struct Config {
155    aliases: HashMap<String, Vec<String>>,
156    bools: HashSet<String>,
157    strings: HashSet<String>,
158    all_bools: bool,
159    defaults: Vec<(String, Value)>,
160    unknown: Option<UnknownFn>,
161}
162
163impl Config {
164    fn new(opts: &ArgvOptions) -> Self {
165        let mut aliases: HashMap<String, Vec<String>> = HashMap::new();
166        for (key, names) in &opts.alias {
167            let group: Vec<&String> = std::iter::once(key).chain(names.iter()).collect();
168            for name in &group {
169                let entry = aliases.entry((*name).clone()).or_default();
170                for other in &group {
171                    if other != name && !entry.contains(other) {
172                        entry.push((*other).clone());
173                    }
174                }
175            }
176        }
177        let mut bools = HashSet::new();
178        if let Booleans::Keys(keys) = &opts.boolean {
179            bools.extend(keys.iter().filter(|k| !k.is_empty()).cloned());
180        }
181        let mut strings = HashSet::new();
182        for key in opts.string.iter().filter(|k| !k.is_empty()) {
183            strings.insert(key.clone());
184            if let Some(names) = aliases.get(key) {
185                strings.extend(names.iter().cloned());
186            }
187        }
188        Self {
189            aliases,
190            bools,
191            strings,
192            all_bools: opts.boolean == Booleans::All,
193            defaults: opts.default.clone(),
194            unknown: opts.unknown,
195        }
196    }
197
198    fn names(&self, key: &str) -> Vec<String> {
199        let mut names = vec![key.to_string()];
200        if let Some(others) = self.aliases.get(key) {
201            names.extend(others.iter().cloned());
202        }
203        names
204    }
205
206    fn alias_is_boolean(&self, key: &str) -> bool {
207        self.aliases
208            .get(key)
209            .is_some_and(|names| names.iter().any(|n| self.bools.contains(n)))
210    }
211
212    fn is_declared(&self, key: &str, arg: &str) -> bool {
213        (self.all_bools && arg.starts_with("--") && !arg.contains('='))
214            || self.strings.contains(key)
215            || self.bools.contains(key)
216            || self.aliases.contains_key(key)
217    }
218
219    fn empty_value(&self, key: &str) -> Value {
220        if self.strings.contains(key) {
221            Value::from("")
222        } else {
223            Value::Bool(true)
224        }
225    }
226}
227
228fn has_path(obj: &Map<String, Value>, keys: &[&str]) -> bool {
229    let mut node = obj;
230    let (last, parents) = keys.split_last().expect("at least one key");
231    for key in parents {
232        match node.get(*key) {
233            Some(Value::Object(child)) => node = child,
234            _ => return false,
235        }
236    }
237    node.contains_key(*last)
238}
239
240fn assign_path(obj: &mut Map<String, Value>, keys: &[&str], value: Value, config: &Config) {
241    if keys.iter().any(|k| UNSAFE_KEYS.contains(k)) {
242        return;
243    }
244    let (last, parents) = keys.split_last().expect("at least one key");
245    let mut node = obj;
246    for key in parents {
247        let slot = node
248            .entry((*key).to_string())
249            .or_insert_with(|| Value::Object(Map::new()));
250        if slot.is_null() {
251            *slot = Value::Object(Map::new());
252        }
253        match slot {
254            Value::Object(child) => node = child,
255            _ => return,
256        }
257    }
258    let replace = config.bools.contains(*last);
259    match node.get_mut(*last) {
260        None => {
261            node.insert((*last).to_string(), value);
262        }
263        Some(existing) if replace || existing.is_boolean() || existing.is_null() => {
264            *existing = value;
265        }
266        Some(Value::Array(items)) => items.push(value),
267        Some(existing) => {
268            let previous = existing.take();
269            *existing = Value::Array(vec![previous, value]);
270        }
271    }
272}
273
274struct Parser<'a> {
275    config: &'a Config,
276    argv: Map<String, Value>,
277    positionals: Vec<Value>,
278}
279
280impl Parser<'_> {
281    fn set_arg(&mut self, key: &str, raw: Value, arg: Option<&str>) {
282        if let (Some(arg), Some(unknown)) = (arg, self.config.unknown) {
283            if !self.config.is_declared(key, arg) && !unknown(arg) {
284                return;
285            }
286        }
287        let value = match raw {
288            Value::String(s) if !self.config.strings.contains(key) && is_number(&s) => {
289                to_number(&s)
290            }
291            other => other,
292        };
293        for name in self.config.names(key) {
294            let path: Vec<&str> = name.split('.').collect();
295            assign_path(&mut self.argv, &path, value.clone(), self.config);
296        }
297    }
298
299    fn takes_value(&self, key: &str, next: Option<&str>, arg: &str) -> bool {
300        next.is_some_and(|n| !FLAG_LIKE_RE.is_match(n))
301            && !self.config.bools.contains(key)
302            && !(self.config.all_bools && arg.starts_with("--"))
303            && !self.config.alias_is_boolean(key)
304    }
305
306    fn set_trailing_flag(&mut self, key: &str, next: Option<&str>, arg: &str) -> usize {
307        if self.takes_value(key, next, arg) {
308            self.set_arg(key, Value::from(next.unwrap_or_default()), Some(arg));
309            return 1;
310        }
311        if let Some(word @ ("true" | "false")) = next {
312            self.set_arg(key, Value::Bool(word == "true"), Some(arg));
313            return 1;
314        }
315        let empty = self.config.empty_value(key);
316        self.set_arg(key, empty, Some(arg));
317        0
318    }
319
320    fn long_flag(&mut self, arg: &str, next: Option<&str>) -> usize {
321        let body = &arg[2..];
322        if let Some((key, raw)) = body.split_once('=').filter(|(k, _)| !k.is_empty()) {
323            let value = if self.config.bools.contains(key) {
324                Value::Bool(raw != "false")
325            } else {
326                Value::from(raw)
327            };
328            self.set_arg(key, value, Some(arg));
329            return 0;
330        }
331        if let Some(key) = body.strip_prefix("no-").filter(|k| !k.is_empty()) {
332            self.set_arg(key, Value::Bool(false), Some(arg));
333            return 0;
334        }
335        self.set_trailing_flag(body, next, arg)
336    }
337
338    /// Handles all letters of a short group but the last; `true` when the
339    /// rest of the group was consumed as a value.
340    fn short_letters(&mut self, arg: &str) -> bool {
341        let chars: Vec<char> = arg.chars().collect();
342        let letters = &chars[1..chars.len() - 1];
343        for (j, &letter) in letters.iter().enumerate() {
344            let key = letter.to_string();
345            let rest: String = chars[j + 2..].iter().collect();
346            let alpha = letter.is_ascii_alphabetic();
347            if rest == "-" {
348                self.set_arg(&key, Value::from(rest), Some(arg));
349            } else if alpha && rest.starts_with('=') {
350                self.set_arg(&key, Value::from(&rest[1..]), Some(arg));
351                return true;
352            } else if (alpha && NUMERIC_TAIL_RE.is_match(&rest))
353                || letters
354                    .get(j + 1)
355                    .is_some_and(|c| !(c.is_ascii_alphanumeric() || *c == '_'))
356            {
357                self.set_arg(&key, Value::from(rest), Some(arg));
358                return true;
359            } else {
360                let empty = self.config.empty_value(&key);
361                self.set_arg(&key, empty, Some(arg));
362            }
363        }
364        false
365    }
366
367    fn short_group(&mut self, arg: &str, next: Option<&str>) -> usize {
368        if self.short_letters(arg) {
369            return 0;
370        }
371        let key = arg.chars().last().map(String::from).unwrap_or_default();
372        if key == "-" {
373            return 0;
374        }
375        self.set_trailing_flag(&key, next, arg)
376    }
377
378    fn push_positional(&mut self, arg: &str) {
379        if let Some(unknown) = self.config.unknown {
380            if !unknown(arg) {
381                return;
382            }
383        }
384        let value = if self.config.strings.contains("_") || !is_number(arg) {
385            Value::from(arg)
386        } else {
387            to_number(arg)
388        };
389        self.positionals.push(value);
390    }
391}
392
393/// Parse `args` like the `minimist` package.
394pub fn minimist<S: AsRef<str>>(args: &[S], opts: &ArgvOptions) -> Value {
395    let config = Config::new(opts);
396    let mut parser = Parser {
397        config: &config,
398        argv: Map::new(),
399        positionals: Vec::new(),
400    };
401    let mut bools: Vec<&String> = config.bools.iter().collect();
402    bools.sort();
403    for key in bools {
404        let fallback = config
405            .defaults
406            .iter()
407            .find(|(k, _)| k == key)
408            .map_or(Value::Bool(false), |(_, v)| v.clone());
409        parser.set_arg(key, fallback, None);
410    }
411
412    let all: Vec<&str> = args.iter().map(AsRef::as_ref).collect();
413    let (list, after) = match all.iter().position(|a| *a == "--") {
414        Some(i) => (&all[..i], &all[i + 1..]),
415        None => (&all[..], &all[all.len()..]),
416    };
417
418    let mut i = 0;
419    while i < list.len() {
420        let arg = list[i];
421        let next = list.get(i + 1).copied();
422        if arg.len() > 2 && arg.starts_with("--") {
423            i += parser.long_flag(arg, next);
424        } else if arg.len() > 1 && arg.starts_with('-') && !arg[1..].starts_with('-') {
425            i += parser.short_group(arg, next);
426        } else {
427            parser.push_positional(arg);
428            if opts.stop_early {
429                parser
430                    .positionals
431                    .extend(list[i + 1..].iter().map(|a| Value::from(*a)));
432                break;
433            }
434        }
435        i += 1;
436    }
437
438    for (key, value) in &config.defaults {
439        let path: Vec<&str> = key.split('.').collect();
440        if has_path(&parser.argv, &path) {
441            continue;
442        }
443        for name in config.names(key) {
444            let path: Vec<&str> = name.split('.').collect();
445            assign_path(&mut parser.argv, &path, value.clone(), &config);
446        }
447    }
448
449    let rest: Vec<Value> = after.iter().map(|a| Value::from(*a)).collect();
450    let mut argv = parser.argv;
451    let mut positionals = parser.positionals;
452    if opts.double_dash {
453        argv.insert("--".into(), Value::Array(rest));
454    } else {
455        positionals.extend(rest);
456    }
457    argv.insert("_".into(), Value::Array(positionals));
458    Value::Object(argv)
459}
460
461fn parse_bool_value(value: Value) -> Value {
462    match value {
463        Value::String(s) => match s.as_str() {
464            "true" => Value::Bool(true),
465            "false" => Value::Bool(false),
466            _ => Value::String(s),
467        },
468        other => other,
469    }
470}
471
472/// zx `parseArgv(args, opts, defs)`: run [`minimist`] and merge the result
473/// into `defs`, optionally camel-casing keys and parsing boolean strings.
474pub fn parse_argv<S: AsRef<str>>(args: &[S], opts: &ArgvOptions, defs: Option<Value>) -> Value {
475    let mut out = match defs {
476        Some(Value::Object(map)) => map,
477        _ => Map::new(),
478    };
479    if let Value::Object(parsed) = minimist(args, opts) {
480        for (key, value) in parsed {
481            if key == "_" || key == "--" {
482                out.insert(key, value);
483                continue;
484            }
485            let key = if opts.camel_case {
486                to_camel_case(&key)
487            } else {
488                key
489            };
490            let value = if opts.parse_boolean {
491                parse_bool_value(value)
492            } else {
493                value
494            };
495            out.insert(key, value);
496        }
497    }
498    Value::Object(out)
499}
500
501/// Parse the current process arguments (without the program name).
502pub fn argv() -> Value {
503    let args: Vec<String> = std::env::args().skip(1).collect();
504    parse_argv(&args, &ArgvOptions::default(), None)
505}