IMprove errror reporting

src/driver/argument_parser.hpp

@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <functional>
 #include <iostream>
+#include <list>
 #include <set>
 #include <string>
 #include <sstream>
@@ -17,6 +18,7 @@
 #include <migraphx/type_name.hpp>
 #include <migraphx/functional.hpp>
 #include <migraphx/stringutils.hpp>
+#include <migraphx/ranges.hpp>
 #include <migraphx/rank.hpp>
 
 #ifndef _WIN32
@@ -81,6 +83,13 @@ inline std::ostream& operator<<(std::ostream& os, const color& c)
     return os;
 }
 
+inline std::string colorize(color c, const std::string& s)
+{
+    std::stringstream ss;
+    ss << c << s << color::reset;
+    return ss.str();
+}
+
 template <class T>
 struct type_name
 {
@@ -248,9 +257,27 @@ struct argument_parser
         });
     }
 
+    template<class F>
+    argument* find_argument(F f)
+    {
+        auto it = std::find_if(arguments.begin(), arguments.end(), f);
+        if (it == arguments.end())
+            return nullptr;
+        return std::addressof(*it);
+    }
+
     MIGRAPHX_DRIVER_STATIC auto show_help(const std::string& msg = "")
     {
         return do_action([=](auto& self) {
+            argument* input_argument = self.find_argument([](const auto& arg) { return arg.flags.empty(); });
+            std::cout << color::fg_yellow << "USAGE:" << color::reset << std::endl;
+            std::cout << "    " << self.exe_name << " <options> ";
+            if (input_argument)
+                std::cout << input_argument->metavar;
+            std::cout << std::endl;
+            std::cout << std::endl;
+            if (self.find_argument([](const auto& arg) { return arg.nargs == 0; }))
+            {
                 std::cout << color::fg_yellow << "FLAGS:" << color::reset << std::endl;
                 std::cout << std::endl;
                 for(auto&& arg : self.arguments)
@@ -282,6 +309,9 @@ struct argument_parser
                     std::cout << arg.help << std::endl;
                 }
                 std::cout << std::endl;
+            }
+            if (self.find_argument([](const auto& arg) { return arg.nargs != 0; }))
+            {
                 std::cout << color::fg_yellow << "OPTIONS:" << color::reset << std::endl;
                 for(auto&& arg : self.arguments)
                 {
@@ -312,6 +342,7 @@ struct argument_parser
                     std::cout << "        " << arg.help << std::endl;
                 }
                 std::cout << std::endl;
+            }
             if(not msg.empty())
                 std::cout << msg << std::endl;
         });
@@ -345,6 +376,44 @@ struct argument_parser
         };
     }
 
+    bool run_action(const argument& arg, const std::string& flag, const std::vector<std::string>& inputs)
+    {
+        std::string msg = "";
+        try
+        {
+            return arg.action(*this, inputs);
+        }
+        catch(const std::exception& e)
+        {
+            msg = e.what();
+        }
+        catch(...)
+        {
+            msg = "unknown exception";
+        }
+        auto show_usage = [&] {
+            std::cout << flag;
+            if(not arg.type.empty())
+                std::cout << " [" << arg.type << "]";
+        };
+        std::cout << color::fg_red << color::bold << "error: " << color::reset;
+        std::cout << "Invalid input to '" << color::fg_yellow;
+        show_usage();
+        std::cout << color::reset << "'" << std::endl;
+        std::cout << "       " << msg << std::endl;
+        std::cout << std::endl;
+        std::cout << color::fg_yellow << "USAGE:" << color::reset << std::endl;
+        std::cout << "    " << exe_name << " ";
+        show_usage();
+        std::cout << std::endl;
+        if (find_argument([](const auto& a) { return contains(a.flags, "--help"); }))
+        {
+            std::cout << std::endl;
+            std::cout << "For more information try '" << color::fg_green << "--help" << color::reset << "'" << std::endl;
+        }
+        return true;
+    }
+
     bool parse(std::vector<std::string> args)
     {
         std::unordered_map<std::string, unsigned> keywords;
@@ -364,7 +433,7 @@ struct argument_parser
             {
                 if(arg_map.count(flag) > 0)
                 {
-                    if(arg.action(*this, arg_map[flag]))
+                    if(run_action(arg, flag, arg_map[flag]))
                         return true;
                 }
             }
@@ -372,6 +441,16 @@ struct argument_parser
         return false;
     }
 
+    void set_exe_name(const std::string& s)
+    {
+        exe_name = s;
+    }
+
+    const std::string& get_exe_name() const
+    {
+        return exe_name;
+    }
+
     using string_map = std::unordered_map<std::string, std::vector<std::string>>;
     template <class IsKeyword>
     static string_map generic_parse(std::vector<std::string> as, IsKeyword is_keyword)
@@ -406,7 +485,8 @@ struct argument_parser
     }
 
     private:
-    std::vector<argument> arguments;
+    std::list<argument> arguments;
+    std::string exe_name = "";
 };
 
 } // namespace MIGRAPHX_INLINE_NS

src/driver/command.hpp

@@ -18,7 +18,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 inline auto& get_commands()
 {
     // NOLINTNEXTLINE
-    static std::unordered_map<std::string, std::function<void(std::vector<std::string> args)>> m;
+    static std::unordered_map<std::string, std::function<void(const std::string& exe_name, std::vector<std::string> args)>> m;
     return m;
 }
 
@@ -42,10 +42,11 @@ const std::string& command_name()
 }
 
 template <class T>
-void run_command(std::vector<std::string> args, bool add_help = false)
+void run_command(const std::string& exe_name, std::vector<std::string> args, bool add_help = false)
 {
     T x;
     argument_parser ap;
+    ap.set_exe_name(exe_name + " " + command_name<T>());
     if(add_help)
         ap(nullptr, {"-h", "--help"}, ap.help("Show help"), ap.show_help());
     x.parse(ap);
@@ -58,7 +59,7 @@ template <class T>
 int auto_register_command()
 {
     auto& m              = get_commands();
-    m[command_name<T>()] = [](std::vector<std::string> args) { run_command<T>(args, true); };
+    m[command_name<T>()] = [](const std::string& exe_name, std::vector<std::string> args) { run_command<T>(exe_name, args, true); };
     return 0;
 }
 

src/driver/main.cpp

@@ -552,26 +552,41 @@ struct onnx : command<onnx>
 
 struct main_command
 {
-    static std::string get_command_help(const std::string& title = "Commands:")
+    static std::string get_command_help(const std::string& title = colorize(color::fg_yellow, "COMMANDS:"))
     {
         std::string result = title + "\n";
-        return std::accumulate(get_commands().begin(),
-                               get_commands().end(),
+        std::vector<std::string> commands(get_commands().size());
+        std::transform(get_commands().begin(), get_commands().end(), commands.begin(), [](const auto& p) {
+            return colorize(color::fg_green, p.first);
+        });
+        std::sort(commands.begin(), commands.end());
+        return std::accumulate(commands.begin(),
+                               commands.end(),
                                result,
-                               [](auto r, auto&& p) { return r + "    " + p.first + "\n"; });
+                               [](auto r, auto&& s) { return r + "    " + s + "\n"; });
     }
     void parse(argument_parser& ap)
     {
         std::string version_str = "MIGraphX Version: " + std::to_string(MIGRAPHX_VERSION_MAJOR) +
                                   "." + std::to_string(MIGRAPHX_VERSION_MINOR);
+        ap(unused, {}, ap.metavar("<command>"), ap.set_value(true));                                  
         ap(nullptr, {"-h", "--help"}, ap.help("Show help"), ap.show_help(get_command_help()));
         ap(nullptr,
            {"-v", "--version"},
            ap.help("Show MIGraphX version"),
            ap.show_help(version_str));
+
+        // Trim command off of exe name
+        ap.set_exe_name(ap.get_exe_name().substr(0, ap.get_exe_name().size() - 5));
     }
 
-    void run() { std::cout << get_command_help("Missing command:") << std::endl; }
+    bool unused = false;
+
+    void run()
+    { 
+        std::cout << color::fg_red << color::bold << "error: " << color::reset;
+        std::cout << get_command_help("missing command:") << std::endl; 
+    }
 };
 
 } // namespace MIGRAPHX_INLINE_NS
@@ -593,11 +608,11 @@ int main(int argc, const char* argv[])
     auto cmd = args.front();
     if(m.count(cmd) > 0)
     {
-        m.at(cmd)({args.begin() + 1, args.end()});
+        m.at(cmd)(argv[0], {args.begin() + 1, args.end()});
     }
     else
     {
-        run_command<main_command>(args);
+        run_command<main_command>(argv[0], args);
     }
 
     return 0;