commander/src/commander.cc

#include <iostream>
#include <vector>
#include <map>
#include <memory>
#include <string>
#include <sstream>
#include <stdarg.h>
#include "commander.h"

#define DEFAULT_MESSAGE_MAX_LEN (30)
#define DEFAULT_ERROR_EXIT_CODE (1)

class CommandImpl;
class OptionImpl;
class ArgumentImpl;

using namespace commander;

void assert(bool condition, const char *fmt, ...)
{
	if (!condition)
	{
		va_list args;
		va_start(args, fmt);
		vfprintf(stderr, fmt, args);
		va_end(args);
		fprintf(stderr, "\n");
		exit(DEFAULT_ERROR_EXIT_CODE);
	}
}

class CastClass
{
public:
	void *m_ptr;
};

#define CMD (static_cast<CommandImpl *>(m_ptr))
#define OPT (static_cast<OptionImpl *>(m_ptr))
#define ARG (static_cast<ArgumentImpl *>(m_ptr))
#define CAST_TO(clazz, ptr) (static_cast<clazz *>(((CastClass *)(ptr))->m_ptr))
inline CommandImpl *CMD_EX(Command &cmd) { return CAST_TO(CommandImpl, &cmd); }
inline CommandImpl *CMD_EX(Command *cmd) { return CAST_TO(CommandImpl, cmd); }
inline OptionImpl *OPT_EX(Option &opt) { return CAST_TO(OptionImpl, &opt); }
inline OptionImpl *OPT_EX(Option *opt) { return CAST_TO(OptionImpl, opt); }
inline ArgumentImpl *ARG_EX(Argument &arg) { return CAST_TO(ArgumentImpl, &arg); }
inline ArgumentImpl *ARG_EX(Argument *arg) { return CAST_TO(ArgumentImpl, arg); }

class OptionImpl
{
public:
	OptionImpl(const std::string &name) : m_name(name) {}

	inline const std::string &name() const { return m_name; }

	inline void shortcut(const std::string &shortcut) { m_shortcut = shortcut; }
	inline const std::string &shortcut() const { return m_shortcut; }

	inline void description(const std::string &description) { m_description = description; }
	inline const std::string &description() const { return m_description; }

	inline void required(bool required) { m_required = required; }
	inline bool required() const { return m_required; }

	void boolean(bool boolean) { m_boolean = boolean; }
	bool boolean() const { return m_boolean; }

	void valueName(const std::string &valueName) { m_valueName = valueName; }
	const std::string &valueName() const { return m_valueName; }

	std::string helpMessage(int maxlen = DEFAULT_MESSAGE_MAX_LEN)
	{
		std::stringstream ss;
		if (m_shortcut.size())
			ss << "-" << m_shortcut << ",";
		ss << "--" << m_name;
		if (!m_boolean)
		{
			ss << " ";
			if (m_required)
				ss << "<";
			else
				ss << "[";

			if (m_valueName.size())
				ss << m_valueName;
			else
				ss << m_name;

			if (m_required)
				ss << ">";
			else
				ss << "]";
		}
		std::string result = ss.str();

		if (m_description.size())
		{
			std::stringstream pad;
			for (size_t i = result.size(); i < maxlen; i++)
				pad << " ";
			pad << m_description;
			result += pad.str();
		}
		return result;
	}

private:
	std::string m_name;
	std::string m_shortcut;
	std::string m_description;
	std::string m_valueName;
	bool m_required = false;
	bool m_boolean = false;
};

class ArgumentImpl
{
public:
	inline ArgumentImpl(const std::string &name) : m_name(name) {}

	inline const std::string &name() const { return m_name; }

	inline void description(const std::string &description) { m_description = description; }
	inline const std::string &description() const { return m_description; }

	inline void required(bool required) { m_required = required; }
	inline bool required() const { return m_required; }

	std::string helpMessage(int maxlen = DEFAULT_MESSAGE_MAX_LEN)
	{
		std::string result = m_name;

		if (m_description.size())
		{
			std::stringstream pad;
			for (size_t i = result.size(); i < maxlen; i++)
				pad << " ";
			pad << m_description;
			result += pad.str();
		}
		return result;
	}

private:
	std::string m_name;
	std::string m_description;
	bool m_required = false;
};

class CommandImpl
{
public:
	CommandImpl(const std::string &name) : m_name(name)
	{
		option("help").shortcut("h").boolean().description("Print this help message");
	}

	const std::string &name() const { return m_name; }

	inline void execute(ExecuteCallback callback) { m_excute = callback; }
	inline ExecuteCallback execute() const { return m_excute; }

	inline void description(const std::string &description) { m_description = description; }
	inline const std::string &description() const { return m_description; }

	Command &command(const std::string &name)
	{
		if (m_commands.find(name) == m_commands.end())
		{
			m_commands.emplace(name, std::make_shared<Command>(name));
		}
		return *m_commands[name];
	}

	Option &option(const std::string &name)
	{
		if (m_options.find(name) == m_options.end())
			m_options.emplace(name, std::make_shared<Option>(name));
		return *m_options[name];
	}

	Argument &argument(const std::string &name)
	{
		if (m_arguments.find(name) == m_arguments.end())
			m_arguments.emplace(name, std::make_shared<Argument>(name));
		return *m_arguments[name];
	}

	void setParent(const CommandImpl &parent)
	{
		for (auto it = parent.m_parent_commands.begin(); it != parent.m_parent_commands.end(); ++it)
		{
			m_parent_commands.push_back(*it);
		}
		m_parent_commands.push_back(parent.m_name);
	}

	std::string usage()
	{
		std::stringstream ss;
		ss << "Usage: ";
		// 名称
		for (auto it = m_parent_commands.begin(); it != m_parent_commands.end(); ++it)
		{
			ss << *it << " ";
		}
		ss << m_name;

		// 命令
		if (m_commands.size())
			ss << " [Command]";

		// 选项
		if (m_options.size())
		{
			if (containsRequiredOptions())
				ss << " <Options>";
			else
				ss << " [Options]";
		}

		// 参数
		for (auto it = m_arguments.begin(); it != m_arguments.end(); ++it)
		{
			ArgumentImpl *impl = ARG_EX(it->second.get());
			if (impl->required())
				ss << " <" << impl->name() << ">";
			else
				ss << " [" << impl->name() << "]";
		}
		ss << std::endl;

		// 描述
		if (m_description.size())
		{
			ss << std::endl;
			ss << m_description << std::endl;
		}

		// 命令定义
		if (m_commands.size())
		{
			ss << std::endl;
			ss << "Commands:" << std::endl;
			for (auto it = m_commands.begin(); it != m_commands.end(); ++it)
			{
				ss << "  " << CMD_EX(it->second.get())->helpMessage() << std::endl;
			}
		}

		// 选项定义
		if (m_options.size())
		{
			ss << std::endl;
			ss << "Options:" << std::endl;
			for (auto it = m_options.begin(); it != m_options.end(); ++it)
			{
				ss << "  " << OPT_EX(it->second.get())->helpMessage() << std::endl;
			}
		}

		// 参数
		if (m_arguments.size())
		{
			ss << std::endl;
			ss << "Arguments:" << std::endl;
			for (auto it = m_arguments.begin(); it != m_arguments.end(); ++it)
			{
				ss << "  " << ARG_EX(it->second.get())->helpMessage() << std::endl;
			}
		}

		return ss.str();
	}

	std::string helpMessage(int maxLen = DEFAULT_MESSAGE_MAX_LEN)
	{
		std::string result = m_name;
		if (m_description.size())
		{
			std::stringstream pad;
			for (size_t i = result.size(); i < maxLen; i++)
			{
				pad << " ";
			}
			pad << m_description;
			result += pad.str();
		}
		return result;
	}

	std::shared_ptr<Option> getOption(const std::string &name, bool isShortcut)
	{
		for (auto &it : m_options)
		{
			OptionImpl *impl = OPT_EX(it.second.get());
			if (isShortcut)
			{
				if (impl->shortcut() == name)
					return it.second;
			}
			else if (impl->name() == name)
				return it.second;
		}
		return nullptr;
	}

	std::shared_ptr<Command> getCommand(const std::string &name)
	{
		for (auto &it : m_commands)
		{
			CommandImpl *impl = CMD_EX(it.second.get());
			if (impl->name() == name)
				return it.second;
		}
		return nullptr;
	}

	std::shared_ptr<Argument> getArgument(int index)
	{
		if (index < 0 || index >= m_arguments.size())
			return nullptr;
		int i = 0;
		for (auto &it : m_arguments)
		{
			if (i == index)
				return it.second;
			++i;
		}
		return nullptr;
	}

	void validate(const std::map<std::string, std::string> &options, const std::map<std::string, std::string> &arguments)
	{
		// 校验选项
		for (auto &it : m_options)
		{
			OptionImpl *impl = OPT_EX(it.second.get());
			auto checkedOption = options.find(impl->name());
			if (impl->required())
				assert(checkedOption != options.end() && checkedOption->second.size(), "Option %s is required", impl->name().c_str());
			if (checkedOption == options.end())
				continue;
			// 其他校验
		}
		// 校验参数
		for (auto &it : m_arguments)
		{
			ArgumentImpl *impl = ARG_EX(it.second.get());
			auto checkedArgument = arguments.find(impl->name());
			if (impl->required())
				assert(checkedArgument != arguments.end() && checkedArgument->second.size(), "Argument %s is required", impl->name().c_str());
			if (checkedArgument == arguments.end())
				continue;
			// 其他校验
		}
	}

private:
	bool containsRequiredArguments()
	{
		for (auto it = m_arguments.begin(); it != m_arguments.end(); ++it)
		{
			if (ARG_EX(it->second.get())->required())
				return true;
		}
		return false;
	}

	bool containsRequiredOptions()
	{
		for (auto it = m_options.begin(); it != m_options.end(); ++it)
		{
			if (OPT_EX(it->second.get())->required())
				return true;
		}
		return false;
	}

private:
	std::string m_name;
	std::vector<std::string> m_parent_commands;
	std::string m_description;
	ExecuteCallback m_excute;
	std::map<std::string, std::shared_ptr<Command>> m_commands;
	std::map<std::string, std::shared_ptr<Option>> m_options;
	std::map<std::string, std::shared_ptr<Argument>> m_arguments;
};

class CommandParser
{
public:
	inline CommandParser(int argc, char **argv) : m_argc(argc), m_argv(argv) {}

	void printUsage(CommandImpl *cmd)
	{
		std::cout << cmd->usage() << std::endl;
		exit(1);
	}

	bool containsHelp()
	{
		for (int i = 0; i < m_argc; i++)
		{
			std::string arg = m_argv[i];
			if (arg == "--help" || arg == "-h")
				return true;
		}
		return false;
	}

	int parse(Command *command)
	{
		std::map<std::string, std::string> options = {};
		std::map<std::string, std::string> arguments = {};
		CommandImpl *cmd = CMD_EX(command);
		bool hasHelp = containsHelp();

		for (int i = 0; i < m_argc; i++)
		{
			std::string arg = m_argv[i];
			// 选项
			if (arg[0] == '-')
			{
				// 如果存在帮助，则直接打印并退出
				if (hasHelp)
					printUsage(cmd);

				std::string name;
				std::string value;
				OptionImpl *impl = nullptr;
				if (arg[1] == '-')
				{
					auto eqIndex = arg.find("=");
					if (eqIndex != std::string::npos)
					{
						name = arg.substr(2, eqIndex - 2);
						value = arg.substr(eqIndex + 1);
						std::cout << name << " " << value << std::endl;
						auto option = cmd->getOption(name, false);
						assert(option != nullptr, "Unknown option: %s", arg.c_str());
						impl = OPT_EX(option.get());
						if (impl->boolean())
						{
							value = "";
						}
					}
					else
					{
						name = arg.substr(2);
						auto option = cmd->getOption(name, false);
						assert(option != nullptr, "Unknown option: %s", arg.c_str());
						impl = OPT_EX(option.get());
						if (i < m_argc - 1 && !impl->boolean())
						{
							value = m_argv[i + 1];
							++i;
						}
					}
				}
				else
				{
					std::string shortcut = arg.substr(1);
					auto option = cmd->getOption(shortcut, true);
					assert(!!option, "Unknown option %s", arg.c_str());
					impl = OPT_EX(option.get());
					name = impl->name();
					if (i < m_argc - 1 && !impl->boolean())
					{
						value = m_argv[i + 1];
						i++;
					}
				}
				options.emplace(name, value);
				continue;
			}
			// 参数或命令
			else
			{
				// 命令检测
				if (!options.size() && !arguments.size())
				{
					auto cmdPtr = cmd->getCommand(arg);
					if (cmdPtr != nullptr)
					{
						cmd = CMD_EX(cmdPtr.get());
						continue;
					}
				}

				// 如果存在帮助，则直接打印并退出
				if (hasHelp)
					printUsage(cmd);
				// 参数
				auto argPtr = cmd->getArgument(static_cast<int>(arguments.size()));
				assert(argPtr != nullptr, "Unknown argument #%d", arguments.size());
				ArgumentImpl *impl = ARG_EX(argPtr.get());
				arguments.emplace(impl->name(), arg);
			}
		}

		// 参数验证
		cmd->validate(options, arguments);

		// 执行函数
		auto fn = cmd->execute();
		if (!fn)
			printUsage(cmd);
		return fn(CommandArgument(options, arguments));
	}

private:
	int m_argc;
	char **m_argv;
	std::map<std::string, std::string> m_options;
	std::vector<std::string> m_commands;
};

Command::Command(const std::string &name) : m_ptr(new CommandImpl(name)) {}
Command::~Command() { delete CMD; }
Command &Command::description(const std::string &description) { return CMD->description(description), *this; }
Command &Command::execute(ExecuteCallback callback) { return CMD->execute(callback), *this; }
Command &Command::command(const std::string &name)
{
	Command &subCommand = CMD->command(name);
	static_cast<CommandImpl *>(subCommand.m_ptr)->setParent(*CMD);
	return subCommand;
}
Option &Command::option(const std::string &name) { return CMD->option(name); }
Argument &Command::argument(const std::string &name) { return CMD->argument(name); }
std::string Command::usage() { return CMD->usage(); }

Option::Option(const std::string &name) : m_ptr(new OptionImpl(name)) {}
Option::~Option() { delete OPT; }
Option &Option::shortcut(const std::string &shortcut) { return OPT->shortcut(shortcut), *this; }
Option &Option::description(const std::string &description) { return OPT->description(description), *this; }
Option &Option::required(bool required) { return OPT->required(required), *this; }
Option &Option::boolean(bool boolean) { return OPT->boolean(boolean), *this; }
Option &Option::valueName(const std::string &valueName) { return OPT->valueName(valueName), *this; }

Argument::Argument(const std::string &name) : m_ptr(new ArgumentImpl(name)) {}
Argument::~Argument() { delete ARG; }
Argument &Argument::description(const std::string &description) { return ARG->description(description), *this; }
Argument &Argument::required(bool required) { return ARG->required(required), *this; }

void CommandArgument::getValue(const std::map<std::string, std::string> &input, bool &out, const std::string &name, bool defaultValue) { out = input.find(name) != input.end(); }

void CommandArgument::getValue(const std::map<std::string, std::string> &input, int32_t &out, const std::string &name, int32_t defaultValue)
{
	auto it = input.find(name);
	if (it != input.end() && !it->second.empty())
	{
		try
		{
			out = std::stoi(it->second);
		}
		catch (...)
		{
			out = defaultValue;
		}
	}
	else
		out = defaultValue;
}

void CommandArgument::getValue(const std::map<std::string, std::string> &input, int64_t &out, const std::string &name, int64_t defaultValue)
{
	auto it = input.find(name);
	if (it != input.end() && !it->second.empty())
	{
		try
		{
			out = std::stoll(it->second);
		}
		catch (...)
		{
			out = defaultValue;
		}
	}
	else
	{
		out = defaultValue;
	}
}

void CommandArgument::getValue(const std::map<std::string, std::string> &input, float &out, const std::string &name, float defaultValue)
{
	auto it = input.find(name);
	if (it != input.end() && !it->second.empty())
	{
		try
		{
			out = std::stof(it->second);
		}
		catch (...)
		{
			out = defaultValue;
		}
	}
	else
	{
		out = defaultValue;
	}
}

void CommandArgument::getValue(const std::map<std::string, std::string> &input, double &out, const std::string &name, double defaultValue)
{
	auto it = input.find(name);
	if (it != input.end() && !it->second.empty())
	{
		try
		{
			out = std::stod(it->second);
		}
		catch (...)
		{
			out = defaultValue;
		}
	}
	else
	{
		out = defaultValue;
	}
}

void CommandArgument::getValue(const std::map<std::string, std::string> &input, std::string &out, const std::string &name, std::string defaultValue)
{
	auto it = input.find(name);
	if (it != input.end() && !it->second.empty())
		out = it->second;
	else
		out = defaultValue;
}

int commander::parse(Command &command, int argc, char **argv)
{
	CommandParser parser(argc, argv);
	return parser.parse(&command);
}