Program Listing for File cli_wrapper.h

↰ Return to documentation for file (src/common/cli_wrapper.h)

#pragma once

#include "3rd_party/CLI/CLI.hpp"
#include "3rd_party/any_type.h"
#include "3rd_party/yaml-cpp/yaml.h"
#include "common/definitions.h"

#include <iostream>
#include <map>
#include <string>
#include <unordered_set>
#include <vector>

namespace marian {

class Options;

namespace cli {

// Option priority
enum struct OptionPriority : int { DefaultValue = 0, ConfigFile = 1, CommandLine = 2 };

struct CLIOptionTuple {
  CLI::Option *opt;     // a pointer to an option object from CLI11
  Ptr<any_type> var;    // value assigned to the option via command-line
  size_t idx{0};        // order in which the option was created
  OptionPriority priority{cli::OptionPriority::DefaultValue};
};

// Helper tuple for aliases storing the alias name, value, and options to be expanded
struct CLIAliasTuple {
  std::string key;    // alias option name
  std::string value;  // value for the alias option indicating that it should be expanded
  YAML::Node config;  // config with options that the alias adds
};

// The helper class for cli::CLIWrapper handling formatting of options and their descriptions.
class CLIFormatter : public CLI::Formatter {
public:
  CLIFormatter(size_t columnWidth, size_t screenWidth);
  virtual std::string make_option_desc(const CLI::Option*) const override;

private:
  size_t screenWidth_{0};
};

class CLIWrapper {
private:
  // Map with option names and option tuples
  std::unordered_map<std::string, CLIOptionTuple> options_;
  // Counter for created options to keep track of order in which options were created
  size_t counter_{0};
  std::vector<CLIAliasTuple> aliases_;  // List of alias tuples

  Ptr<CLI::App> app_;                   // Command-line argument parser from CLI11

  std::string defaultGroup_{""};        // Name of the default option group
  std::string currentGroup_{""};        // Name of the current option group

  YAML::Node &config_;                  // Reference to the main config object

  // Option for --version flag. This is a special flag and similarly to --help,
  // the key "version" will be not added into the YAML config
  CLI::Option *optVersion_;

  // Extract option name from a comma-separated list of long and short options, e.g. 'help' from
  // '--help,-h'
  std::string keyName(const std::string &args) const;

  // Get names of options passed via command-line
  std::unordered_set<std::string> getParsedOptionNames() const;
  // Get option names in the same order as they are created
  std::vector<std::string> getOrderedOptionNames() const;

  static std::string failureMessage(const CLI::App *app, const CLI::Error &e);

public:

  CLIWrapper(YAML::Node &config,
             const std::string &description = "",
             const std::string &header = "General options",
             const std::string &footer = "",
             size_t columnWidth = 40,
             size_t screenWidth = 0);

  virtual ~CLIWrapper();

  template <typename T>
  CLI::Option* add(const std::string& args, const std::string& help, T val) {
    return addOption<T>(keyName(args),
                        args,
                        help,
                        val,
                        /*defaulted =*/true);
  }

  template <typename T>
  CLI::Option* add(const std::string& args, const std::string& help) {
    return addOption<T>(keyName(args),
                        args,
                        help,
                        T(),
                        /*defaulted =*/false);
  }

  void alias(const std::string &key,
             const std::string &value,
             const std::function<void(YAML::Node &config)> &fun) {
    ABORT_IF(!options_.count(key), "Option '{}' is not defined so alias can not be created", key);
    aliases_.resize(aliases_.size() + 1);
    aliases_.back().key = key;
    aliases_.back().value = value;
    fun(aliases_.back().config);
  }

  std::string switchGroup(std::string name = "");

  // Parse command-line arguments. Handles --help and --version options
  void parse(int argc, char** argv);

  void parseAliases();

  void updateConfig(const YAML::Node &config, cli::OptionPriority priority, const std::string &errorMsg);

  // Get textual YAML representation of the config
  std::string dumpConfig(bool skipUnmodified = false) const;

private:
  template <typename T>
  using EnableIfNumbericOrString = CLI::enable_if_t<!CLI::is_bool<T>::value
                                   && !CLI::is_vector<T>::value, CLI::detail::enabler>;

  template <typename T, EnableIfNumbericOrString<T> = CLI::detail::dummy>
  CLI::Option* addOption(const std::string &key,
                         const std::string &args,
                         const std::string &help,
                         T val,
                         bool defaulted) {
    // add key to YAML
    config_[key] = val;

    // create option tuple
    CLIOptionTuple option;
    option.idx = counter_++;
    option.var = std::make_shared<any_type>(val);

    // callback function collecting a command-line argument
    CLI::callback_t fun = [this, key](CLI::results_t res) {
      options_[key].priority = cli::OptionPriority::CommandLine;
      // get variable associated with the option
      auto& var = options_[key].var->as<T>();
      // store parser result in var
      auto ret = CLI::detail::lexical_cast(res[0], var);
      // update YAML entry
      config_[key] = var;
      return ret;
    };

    auto opt = app_->add_option(args, fun, help, defaulted);
    // set human readable type value: UINT, INT, FLOAT or TEXT
    opt->type_name(CLI::detail::type_name<T>());
    // set option group
    if(!currentGroup_.empty())
      opt->group(currentGroup_);
    // set textual representation of the default value for help message
    if(defaulted) {
      std::stringstream ss;
      ss << val;
      opt->default_str(ss.str());
    }

    // store option tuple
    option.opt = opt;
    options_.insert(std::make_pair(key, option));
    return options_[key].opt;
  }

  template <typename T>
  using EnableIfVector = CLI::enable_if_t<CLI::is_vector<T>::value, CLI::detail::enabler>;

  template <typename T, EnableIfVector<T> = CLI::detail::dummy>
  CLI::Option* addOption(const std::string &key,
                         const std::string &args,
                         const std::string &help,
                         T val,
                         bool defaulted) {
    // add key to YAML
    config_[key] = val;

    // create option tuple
    CLIOptionTuple option;
    option.idx = counter_++;
    option.var = std::make_shared<any_type>(val);

    // callback function collecting command-line arguments
    CLI::callback_t fun = [this, key](CLI::results_t res) {
      options_[key].priority = cli::OptionPriority::CommandLine;
      // get vector variable associated with the option
      auto& vec = options_[key].var->as<T>();
      vec.clear();
      bool ret = true;
      // handle '[]' as an empty vector
      if(res.size() == 1 && res.front() == "[]") {
        ret = true;
      } else {
        // populate the vector with parser results
        for(const auto& a : res) {
          vec.emplace_back();
          ret &= CLI::detail::lexical_cast(a, vec.back());
        }
        ret &= !vec.empty();
      }
      // update YAML entry
      config_[key] = vec;
      return ret;
    };

    auto opt = app_->add_option(args, fun, help);
    // set human readable type value: VECTOR and
    opt->type_name(CLI::detail::type_name<T>());
    // accept unlimited number of arguments
    opt->type_size(-1);
    // set option group
    if(!currentGroup_.empty())
      opt->group(currentGroup_);
    // set textual representation of the default vector values for help message
    if(defaulted)
      opt->default_str(CLI::detail::join(val));

    // store option tuple
    option.opt = opt;
    options_.insert(std::make_pair(key, option));
    return options_[key].opt;
  }

  template <typename T>
  using EnableIfBoolean = CLI::enable_if_t<CLI::is_bool<T>::value, CLI::detail::enabler>;

  template <typename T, EnableIfBoolean<T> = CLI::detail::dummy>
  CLI::Option* addOption(const std::string &key,
                         const std::string &args,
                         const std::string &help,
                         T val,
                         bool defaulted) {
    // add key to YAML
    config_[key] = val;

    // create option tuple
    CLIOptionTuple option;
    option.idx = counter_++;
    option.var = std::make_shared<any_type>(val);

    // callback function setting the flag
    CLI::callback_t fun = [this, key](CLI::results_t res) {
      options_[key].priority = cli::OptionPriority::CommandLine;
      // get parser result, it is safe as boolean options have an implicit value
      auto val = res[0];
      auto ret = true;
      if(val == "true" || val == "on" || val == "yes" || val == "1") {
        options_[key].var->as<T>() = true;
        config_[key] = true;
      } else if(val == "false" || val == "off" || val == "no" || val == "0") {
        options_[key].var->as<T>() = false;
        config_[key] = false;
      } else {
        ret = false;
      }
      return ret;
    };

    auto opt = app_->add_option(args, fun, help, defaulted);
    // set option group
    if(!currentGroup_.empty())
      opt->group(currentGroup_);
    // set textual representation of the default value for help message
    if(defaulted)
      opt->default_str(val ? "true" : "false");
    // allow to use the flag without any argument
    opt->implicit_val("true");

    // store option tuple
    option.opt = opt;
    options_.insert(std::make_pair(key, option));
    return options_[key].opt;
  }
};

}  // namespace cli
}  // namespace marian