For Json flags, if the input does not parse as a json object try to treat it as a string.

PiperOrigin-RevId: 750371763
Change-Id: I5ad593e7e8ca6cbc11734975749ae52824503e0d

Author: laramiel

tensorstore/util/BUILD

@@ -438,6 +438,7 @@ tensorstore_cc_library(
         "@com_github_nlohmann_json//:json",
         "@com_google_absl//absl/flags:marshalling",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
     ],
 )

tensorstore/util/json_absl_flag.h

@@ -21,6 +21,7 @@
 
 #include "absl/flags/marshalling.h"
 #include "absl/status/status.h"
+#include "absl/strings/ascii.h"
 #include "absl/strings/str_format.h"
 #include <nlohmann/json.hpp>
 #include "tensorstore/internal/json_binding/bindable.h"
@@ -58,22 +59,28 @@ struct JsonAbslFlag {
 
   friend bool AbslParseFlag(std::string_view in, JsonAbslFlag* out,
                             std::string* error) {
+    // Whitespace is removed from JSON values: https://www.json.org/
+    in = absl::StripAsciiWhitespace(in);
     if (in.empty()) {
-      // empty string yields a default-constructed flag value.
       out->value = {};
       return true;
     }
 
     ::nlohmann::json j = ::nlohmann::json::parse(in, nullptr, false);
     if (j.is_discarded()) {
-      *error = absl::StrFormat("Failed to parse JSON: '%s'", in);
-      return false;
+      if (in[0] == '"' || in[0] == '{' || in[0] == '[') {
+        *error = absl::StrFormat("Failed to parse JSON: '%s'", in);
+        return false;
+      }
+      j = ::nlohmann::json(std::string(in));
     }
+
     T new_value = {};
     absl::Status status = internal_json_binding::DefaultBinder<>(
         std::true_type{}, internal_json_binding::NoOptions{}, &new_value, &j);
     if (!status.ok()) {
-      *error = absl::StrFormat("Failed to bind JSON: %s", status.message());
+      *error =
+          absl::StrFormat("Failed to parse or bind JSON: %s", status.message());
       return false;
     }
     out->value = std::move(new_value);

tensorstore/util/json_absl_flag_test.cc

@@ -18,31 +18,53 @@
 
 #include <string>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "tensorstore/internal/json_binding/json_binding.h"
+#include "tensorstore/internal/json_binding/json_binding.h"  // IWYU pragma: keep
 #include "tensorstore/kvstore/spec.h"
 
 namespace {
 
 TEST(JsonAbslFlag, IntFlag) {
   // Validate that the default value can roundtrip.
   tensorstore::JsonAbslFlag<int64_t> flag = {};
-  std::string default_value = AbslUnparseFlag(flag);
+  EXPECT_THAT(AbslUnparseFlag(flag), ::testing::StrEq("0"));
 
   std::string error;
-  EXPECT_TRUE(AbslParseFlag(default_value, &flag, &error));
+  EXPECT_TRUE(AbslParseFlag("", &flag, &error));
   EXPECT_TRUE(error.empty());
+
+  EXPECT_TRUE(AbslParseFlag("1", &flag, &error));
+  EXPECT_TRUE(error.empty()) << error;
 }
 
 TEST(JsonAbslFlag, KvStoreSpecFlag) {
   // Validate that the default value can roundtrip.
   tensorstore::JsonAbslFlag<tensorstore::kvstore::Spec> flag = {};
-  std::string default_value = AbslUnparseFlag(flag);
+  EXPECT_THAT(AbslUnparseFlag(flag), ::testing::IsEmpty());
 
+  // Try to parse as an empty value JSON string.
   std::string error;
-  EXPECT_TRUE(AbslParseFlag(default_value, &flag, &error))
-      << "value: " << default_value;
+  EXPECT_TRUE(AbslParseFlag("", &flag, &error));
+  EXPECT_TRUE(error.empty()) << error;
+
+  EXPECT_TRUE(AbslParseFlag("  ", &flag, &error));
   EXPECT_TRUE(error.empty()) << error;
+
+  // Try to parse a bad json object.
+  error.clear();
+  EXPECT_FALSE(AbslParseFlag("{ \"driver\": \"memory\" ", &flag, &error));
+  EXPECT_THAT(error, testing::HasSubstr("Failed to parse JSON"));
+
+  // Try to parse as a json object.
+  error.clear();
+  EXPECT_FALSE(AbslParseFlag("{ \"driver\": \"memory\" }", &flag, &error));
+  EXPECT_THAT(error, testing::HasSubstr("Failed to parse or bind JSON"));
+
+  // Try to parse as a raw string.
+  error.clear();
+  EXPECT_FALSE(AbslParseFlag("memory://", &flag, &error));
+  EXPECT_THAT(error, testing::HasSubstr("Failed to parse or bind JSON"));
 }
 
 }  // namespace