mz_ore/
cli.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Command-line parsing utilities.
17
18use std::ffi::OsString;
19use std::fmt::{self, Display};
20use std::str::FromStr;
21
22use clap::Parser;
23
24/// A help template for use with clap that does not include the name of the
25/// binary or the version in the help output.
26const NO_VERSION_HELP_TEMPLATE: &str = "{about}
27
28USAGE:
29    {usage}
30
31{all-args}";
32
33/// Configures command-line parsing via [`parse_args`].
34#[derive(Debug, Default, Clone)]
35pub struct CliConfig<'a> {
36    /// An optional prefix to apply to the environment variable name for  all
37    /// arguments with an environment variable fallback.
38    //
39    // TODO(benesch): switch to the clap-native `env_prefix` option if that
40    // gets implemented: https://github.com/clap-rs/clap/issues/3221.
41    pub env_prefix: Option<&'a str>,
42    /// Enable clap's built-in `--version` flag.
43    ///
44    /// We disable this by default because most of our binaries are not
45    /// meaningfully versioned.
46    pub enable_version_flag: bool,
47}
48
49/// Parses command-line arguments according to a clap `Parser` after
50/// applying Materialize-specific customizations.
51pub fn parse_args<O>(config: CliConfig) -> O
52where
53    O: Parser,
54{
55    // Construct the prefixed environment variable names for all
56    // environment-enabled arguments, if requested. We have to construct these
57    // names before constructing `clap` below to get the lifetimes to work out.
58    let command = O::command();
59    let arg_envs: Vec<_> = command
60        .get_arguments()
61        .filter_map(|arg| match (config.env_prefix, arg.get_env()) {
62            (Some(prefix), Some(env)) => {
63                let mut prefixed_env = OsString::from(prefix);
64                prefixed_env.push(env);
65                Some((arg.get_id(), prefixed_env))
66            }
67            _ => None,
68        })
69        .collect();
70
71    let mut clap = O::command().args_override_self(true);
72
73    if !config.enable_version_flag {
74        clap = clap.disable_version_flag(true);
75        clap = clap.help_template(NO_VERSION_HELP_TEMPLATE);
76    }
77
78    for (arg, env) in &arg_envs {
79        clap = clap.mut_arg(*arg, |arg| arg.env_os(env));
80    }
81
82    O::from_arg_matches(&clap.get_matches()).unwrap()
83}
84
85/// A command-line argument of the form `KEY=VALUE`.
86#[derive(Debug, Clone)]
87pub struct KeyValueArg<K, V> {
88    /// The key of the command-line argument.
89    pub key: K,
90    /// The value of the command-line argument.
91    pub value: V,
92}
93
94impl<K, V> FromStr for KeyValueArg<K, V>
95where
96    K: FromStr,
97    K::Err: Display,
98    V: FromStr,
99    V::Err: Display,
100{
101    type Err = String;
102
103    fn from_str(s: &str) -> Result<KeyValueArg<K, V>, String> {
104        let mut parts = s.splitn(2, '=');
105        let key = parts.next().expect("always one part");
106        let value = parts
107            .next()
108            .ok_or_else(|| "must have format KEY=VALUE".to_string())?;
109        Ok(KeyValueArg {
110            key: key.parse().map_err(|e| format!("parsing key: {}", e))?,
111            value: value.parse().map_err(|e| format!("parsing value: {}", e))?,
112        })
113    }
114}
115
116/// A command-line argument that defaults to `true`
117/// that can be falsified with `--flag=false`
118// TODO: replace with `SetTrue` when available in clap.
119#[derive(Debug, Clone)]
120pub struct DefaultTrue {
121    /// The value for this flag
122    pub value: bool,
123}
124
125impl Default for DefaultTrue {
126    fn default() -> Self {
127        Self { value: true }
128    }
129}
130
131impl fmt::Display for DefaultTrue {
132    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
133        self.value.fmt(f)
134    }
135}
136
137impl FromStr for DefaultTrue {
138    type Err = std::str::ParseBoolError;
139
140    fn from_str(s: &str) -> Result<Self, Self::Err> {
141        Ok(Self { value: s.parse()? })
142    }
143}
mz_ore/cli.rs

mz_ore/
cli.rs
