argument_parser.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.

#ifndef LBANN_UTILS_ARGUMENT_PARSER_HPP_INCLUDED
#define LBANN_UTILS_ARGUMENT_PARSER_HPP_INCLUDED

#include "lbann/utils/environment_variable.hpp"
#include "lbann/utils/exception.hpp"

#include <clara.hpp>

#include <any>
#include <initializer_list>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>

namespace lbann {
namespace utils {

struct parse_error : std::runtime_error
{
  template <typename T>
  parse_error(T&& what_arg) : std::runtime_error{std::forward<T>(what_arg)}
  {}
}; // parse_error

struct strict_parsing
{
  void handle_error(clara::detail::InternalParseResult result,
                    clara::Parser& parser,
                    std::vector<char const*>& argv);
}; // struct strict_parsing

struct allow_extra_parameters
{
  void handle_error(clara::detail::InternalParseResult result,
                    clara::Parser& parser,
                    std::vector<char const*>& argv);
}; // struct allow_extra_parameters

template <typename ErrorHandler>
class argument_parser : ErrorHandler
{
public:

  template <typename T>
  class readonly_reference
  {
  public:
    readonly_reference(T& val) noexcept : ref_(val) {}
    T const& get() const noexcept { return ref_; }
    operator T const&() const noexcept { return this->get(); }

    template <typename S>
    bool operator==(S const& y) const noexcept
    {
      return this->get() == y;
    }

  private:
    T& ref_;
  }; // class readonly_reference<T>

  struct parse_error : std::runtime_error
  {
    template <typename T>
    parse_error(T&& what_arg) : std::runtime_error{std::forward<T>(what_arg)}
    {}
  };

  struct missing_required_arguments : std::runtime_error
  {
    template <typename Container>
    missing_required_arguments(Container const& missing_args)
      : std::runtime_error{build_what_string_(missing_args)}
    {}

  private:
    template <typename Container>
    std::string build_what_string_(Container const& missing_args)
    {
      std::ostringstream oss;
      oss << "The following required arguments are missing: {";
      for (auto const& x : missing_args)
        oss << " \"" << x << "\"";
      oss << " }";
      return oss.str();
    }
  };


public:

  argument_parser();

  argument_parser(argument_parser const&) = delete;

  argument_parser& operator=(argument_parser const&) = delete;

  argument_parser(argument_parser&&) = default;

  argument_parser& operator=(argument_parser&&) = default;



  readonly_reference<bool>
  add_flag(std::string const& name,
           std::initializer_list<std::string> cli_flags,
           std::string const& description);

  template <typename AccessPolicy>
  readonly_reference<bool>
  add_flag(std::string const& name,
           std::initializer_list<std::string> cli_flags,
           EnvVariable<AccessPolicy> env,
           std::string const& description)
  {
    if (env.exists() && env.template value<bool>())
      return add_flag_impl_(name,
                            std::move(cli_flags),
                            description + "\nENV: {" + env.name() + "}",
                            true);
    else
      return add_flag(name,
                      std::move(cli_flags),
                      description + "\nENV: {" + env.name() + "}");
  }

  template <typename T>
  readonly_reference<T> add_option(std::string const& name,
                                   std::initializer_list<std::string> cli_flags,
                                   std::string const& description,
                                   T default_value = T());

  template <typename T, typename AccessPolicy>
  readonly_reference<T> add_option(std::string const& name,
                                   std::initializer_list<std::string> cli_flags,
                                   EnvVariable<AccessPolicy> env,
                                   std::string const& description,
                                   T default_value = T())
  {
    if (env.exists())
      return add_option(name,
                        std::move(cli_flags),
                        description + "\nENV: {" + env.name() + "}",
                        env.template value<T>());
    else
      return add_option(name,
                        std::move(cli_flags),
                        description + "\nENV: {" + env.name() + "}",
                        std::move(default_value));
  }

  readonly_reference<std::string>
  add_option(std::string const& name,
             std::initializer_list<std::string> cli_flags,
             std::string const& description,
             char const* default_value)
  {
    return add_option(name,
                      std::move(cli_flags),
                      description,
                      std::string(default_value));
  }

  template <typename AccessPolicy>
  readonly_reference<std::string>
  add_option(std::string const& name,
             std::initializer_list<std::string> cli_flags,
             EnvVariable<AccessPolicy> env,
             std::string const& description,
             char const* default_value)
  {
    return add_option(name,
                      cli_flags,
                      std::move(env),
                      description + "\nENV: {" + env.name() + "}",
                      std::string(default_value));
  }

  template <typename T>
  readonly_reference<T> add_argument(std::string const& name,
                                     std::string const& description,
                                     T default_value = T());

  readonly_reference<std::string> add_argument(std::string const& name,
                                               std::string const& description,
                                               char const* default_value)
  {
    return add_argument(name, description, std::string(default_value));
  }

  template <typename T>
  readonly_reference<T> add_required_argument(std::string const& name,
                                              std::string const& description);

  void clear() noexcept;



  void parse(int argc, char const* const argv[]);

  void parse_no_finalize(int argc, char const* const argv[]);

  void finalize() const;



  std::string get_exe_name() const noexcept;

  bool option_is_defined(std::string const& option_name) const;

  bool help_requested() const;

  template <typename T>
  T const& get(std::string const& option_name) const;



  void print_help(std::ostream& stream) const;


private:
  void init() noexcept;

  readonly_reference<bool>
  add_flag_impl_(std::string const& name,
                 std::initializer_list<std::string> cli_flags,
                 std::string const& description,
                 bool default_value);

private:
  std::unordered_map<std::string, std::any> params_;
  std::unordered_set<std::string> required_;
  clara::Parser parser_;
};

template <typename ErrorHandler>
inline bool argument_parser<ErrorHandler>::option_is_defined(
  std::string const& option_name) const
{
  return params_.count(option_name);
}

template <typename ErrorHandler>
template <typename T>
inline T const&
argument_parser<ErrorHandler>::get(std::string const& option_name) const
{
  if (!option_is_defined(option_name)) {
    LBANN_ERROR("Invalid option: ", option_name);
  }
  return std::any_cast<T const&>(params_.at(option_name));
}

template <typename ErrorHandler>
template <typename T>
inline auto argument_parser<ErrorHandler>::add_option(
  std::string const& name,
  std::initializer_list<std::string> cli_flags,
  std::string const& description,
  T default_value) -> readonly_reference<T>
{
  params_[name] = std::move(default_value);
  auto& param_ref = std::any_cast<T&>(params_[name]);
  clara::Opt option(param_ref, name);
  for (auto const& f : cli_flags)
    option[f];
  parser_ |= option(description).optional();
  return param_ref;
}

template <typename ErrorHandler>
template <typename T>
inline auto
argument_parser<ErrorHandler>::add_argument(std::string const& name,
                                            std::string const& description,
                                            T default_value)
  -> readonly_reference<T>
{
  params_[name] = std::move(default_value);
  auto& param_ref = std::any_cast<T&>(params_[name]);
  parser_ |= clara::Arg(param_ref, name)(description).optional();
  return param_ref;
}

template <typename ErrorHandler>
template <typename T>
inline auto argument_parser<ErrorHandler>::add_required_argument(
  std::string const& name,
  std::string const& description) -> readonly_reference<T>
{
  // Add the reference to bind to
  params_[name] = T{};
  auto& param_any = params_[name];
  auto& param_ref = std::any_cast<T&>(param_any);

  required_.insert(name);

  // Make sure the required arguments are all grouped together.
  auto iter = parser_.m_args.cbegin(), invalid = parser_.m_args.cend();
  while (iter != invalid && !iter->isOptional())
    ++iter;

  // Create the argument
  auto ret = parser_.m_args.emplace(
    iter,
    [name, &param_ref, this](std::string const& value) {
      auto result = clara::detail::convertInto(value, param_ref);
      if (result)
        required_.erase(name);
      return result;
    },
    name);
  ret->operator()(description).required();
  return param_ref;
}

template <typename ErrorHandler>
argument_parser<ErrorHandler>::argument_parser()
{
  init();
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::clear() noexcept
{
  std::unordered_map<std::string, std::any>{}.swap(params_);
  std::unordered_set<std::string>{}.swap(required_);
  parser_ = clara::Parser{};
  init();
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::init() noexcept
{
  params_["print help"] = false;
  parser_ |= clara::ExeName();
  parser_ |= clara::Help(std::any_cast<bool&>(params_["print help"]));
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::parse(int argc, char const* const argv[])
{
  parse_no_finalize(argc, argv);
  finalize();
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::parse_no_finalize(int argc,
                                                      char const* const argv[])
{
  std::vector<char const*> newargv(argv, argv + argc);
  auto parse_result =
    parser_.parse(clara::Args(newargv.size(), newargv.data()));

  if (!parse_result)
    this->handle_error(parse_result, parser_, newargv);
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::finalize() const
{
  if (!help_requested() && required_.size())
    throw missing_required_arguments(required_);
}

template <typename ErrorHandler>
auto argument_parser<ErrorHandler>::add_flag(
  std::string const& name,
  std::initializer_list<std::string> cli_flags,
  std::string const& description) -> readonly_reference<bool>
{
  return add_flag_impl_(name, std::move(cli_flags), description, false);
}

template <typename ErrorHandler>
std::string argument_parser<ErrorHandler>::get_exe_name() const noexcept
{
  return parser_.m_exeName.name();
}

template <typename ErrorHandler>
bool argument_parser<ErrorHandler>::help_requested() const
{
  return std::any_cast<bool>(params_.at("print help"));
}

template <typename ErrorHandler>
void argument_parser<ErrorHandler>::print_help(std::ostream& out) const
{
  out << parser_ << std::endl;
}

template <typename ErrorHandler>
auto argument_parser<ErrorHandler>::add_flag_impl_(
  std::string const& name,
  std::initializer_list<std::string> cli_flags,
  std::string const& description,
  bool default_value) -> readonly_reference<bool>
{
  params_[name] = default_value;
  auto& param_ref = std::any_cast<bool&>(params_[name]);
  clara::Opt option(param_ref);
  for (auto const& f : cli_flags)
    option[f];
  parser_ |= option(description).optional();
  return param_ref;
}

} // namespace utils

using default_arg_parser_type = utils::argument_parser<utils::strict_parsing>;

default_arg_parser_type& global_argument_parser();

} // namespace lbann

template <typename ErrorHandler>
std::ostream&
operator<<(std::ostream& os,
           lbann::utils::argument_parser<ErrorHandler> const& parser)
{
  parser.print_help(os);
  return os;
}

#endif /* LBANN_UTILS_ARGUMENT_PARSER_HPP_INCLUDED */