From 9d669764219f7f9aec8098a6eb11d6ecdce2fb43 Mon Sep 17 00:00:00 2001
From: Zhihao Yuan <lichray@gmail.com>
Date: Wed, 13 Nov 2019 02:17:28 -0600
Subject: [PATCH] Print ArgumentParser help with stream insertion

---
 include/argparse.hpp | 81 +++++++++++++++++++++++++++-----------------
 test/main.cpp        |  3 +-
 test/test_help.hpp   | 18 ++++++++++
 3 files changed, 69 insertions(+), 33 deletions(-)
 create mode 100644 test/test_help.hpp

diff --git a/include/argparse.hpp b/include/argparse.hpp
index 67783ad..0b3410a 100644
--- a/include/argparse.hpp
+++ b/include/argparse.hpp
@@ -31,7 +31,6 @@ SOFTWARE.
 #include <algorithm>
 #include <any>
 #include <functional>
-#include <iomanip>
 #include <iostream>
 #include <iterator>
 #include <list>
@@ -75,8 +74,12 @@ template <typename T>
 using enable_if_not_container = std::enable_if_t<!is_container_v<T>, T>;
 } // namespace
 
+class ArgumentParser;
+
 class Argument {
   friend class ArgumentParser;
+  friend auto operator<<(std::ostream &, ArgumentParser const &)
+      -> std::ostream &;
 
 public:
   Argument() = default;
@@ -412,39 +415,53 @@ public:
     throw std::logic_error("No such argument");
   }
 
+  // Print help message
+  friend auto operator<<(std::ostream &stream, const ArgumentParser &parser)
+      -> std::ostream & {
+    if (auto sen = std::ostream::sentry(stream)) {
+      stream.setf(std::ios_base::left);
+      stream << "Usage: " << parser.mProgramName << " [options] ";
+      size_t tLongestArgumentLength = parser.get_length_of_longest_argument();
+
+      for (const auto &argument : parser.mPositionalArguments) {
+        stream << argument->mNames.front() << " ";
+      }
+      stream << "\n\n";
+
+      if (!parser.mPositionalArguments.empty())
+        stream << "Positional arguments:\n";
+
+      for (const auto &mPositionalArgument : parser.mPositionalArguments) {
+        stream.width(tLongestArgumentLength);
+        stream << *mPositionalArgument;
+      }
+
+      if (!parser.mOptionalArguments.empty())
+        stream << (parser.mPositionalArguments.empty() ? "" : "\n")
+               << "Optional arguments:\n";
+
+      for (const auto &mOptionalArgument : parser.mOptionalArguments) {
+        stream.width(tLongestArgumentLength);
+        stream << *mOptionalArgument;
+      }
+    }
+
+    return stream;
+  }
+
+  // Format help message
+  auto help() const -> std::ostringstream {
+    std::ostringstream out;
+    out << *this;
+    return out;
+  }
+
   // Printing the one and only help message
   // I've stuck with a simple message format, nothing fancy.
-  // TODO: support user-defined help and usage messages for the ArgumentParser
   std::string print_help() {
-    std::stringstream stream;
-    stream << std::left;
-    stream << "Usage: " << mProgramName << " [options] ";
-    size_t tLongestArgumentLength = get_length_of_longest_argument();
-
-    for (const auto &argument : mPositionalArguments) {
-      stream << argument->mNames.front() << " ";
-    }
-    stream << "\n\n";
-
-    if (!mPositionalArguments.empty())
-      stream << "Positional arguments:\n";
-
-    for (const auto &mPositionalArgument : mPositionalArguments) {
-      stream.width(tLongestArgumentLength);
-      stream << *mPositionalArgument;
-    }
-
-    if (!mOptionalArguments.empty())
-      stream << (mPositionalArguments.empty() ? "" : "\n")
-             << "Optional arguments:\n";
-
-    for (const auto &mOptionalArgument : mOptionalArguments) {
-      stream.width(tLongestArgumentLength);
-      stream << *mOptionalArgument;
-    }
-
-    std::cout << stream.str();
-    return stream.str();
+    auto out = help();
+    std::cout << out.rdbuf();
+    return out.str();
   }
 
 private:
@@ -507,7 +524,7 @@ private:
   }
 
   // Used by print_help.
-  size_t get_length_of_longest_argument() {
+  size_t get_length_of_longest_argument() const {
     if (mArgumentMap.empty())
       return 0;
     std::vector<size_t> argumentLengths(mArgumentMap.size());
diff --git a/test/main.cpp b/test/main.cpp
index 8d9aa37..03eb404 100644
--- a/test/main.cpp
+++ b/test/main.cpp
@@ -1,6 +1,7 @@
 #define CATCH_CONFIG_MAIN
 #include <iostream>
 #include <argparse.hpp>
+#include <test_help.hpp>
 #include <test_parse_args.hpp>
 #include <test_positional_arguments.hpp>
 #include <test_optional_arguments.hpp>
@@ -10,4 +11,4 @@
 #include <test_parent_parsers.hpp>
 #include <test_invalid_arguments.hpp>
 #include <test_negative_numbers.hpp>
-#include <test_required_arguments.hpp>
\ No newline at end of file
+#include <test_required_arguments.hpp>
diff --git a/test/test_help.hpp b/test/test_help.hpp
new file mode 100644
index 0000000..d547630
--- /dev/null
+++ b/test/test_help.hpp
@@ -0,0 +1,18 @@
+#pragma once
+#include <catch.hpp>
+#include <argparse.hpp>
+
+TEST_CASE("Users can format help message", "[help]") {
+  argparse::ArgumentParser program("test");
+  program.add_argument("input")
+    .help("positional input");
+  program.add_argument("-c")
+    .help("optional input");
+
+  std::ostringstream s;
+  s << program;
+  REQUIRE_FALSE(s.str().empty());
+
+  auto msg = program.help().str();
+  REQUIRE(msg == s.str());
+}