Add size output to tscli list

PiperOrigin-RevId: 700846119
Change-Id: Iae7d96361d72d24e4ac7efd2de0585308f420847

tensorstore/tscli/BUILD

@@ -56,9 +56,11 @@ tensorstore_cc_library(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/flags:parse",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
         "@com_google_re2//:re2",

tensorstore/tscli/kvstore_list.cc

@@ -12,15 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
 #include <iostream>
 #include <optional>
 #include <string>
 #include <string_view>
 #include <vector>
 
+#include "absl/functional/function_ref.h"
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "re2/re2.h"
 #include "tensorstore/context.h"
 #include "tensorstore/internal/json_binding/std_optional.h"  // IWYU pragma: keep
@@ -47,10 +53,63 @@ std::string_view NonGlobPrefix(std::string_view glob) {
   return glob.substr(0, first_glob_char);
 }
 
+std::string FormatBytesHumanReadable(int64_t num_bytes) {
+  if (num_bytes < 0) return "";
+
+  // Special case for bytes.
+  if (num_bytes < 1024) {
+    return absl::StrFormat("%dB", num_bytes);
+  }
+
+  // Discard the least significant bits when finding SI prefix.
+  static const char units[] = "KMGTPE";
+  const char* unit = units;
+  while (num_bytes >= 1024 * 1024) {
+    num_bytes /= static_cast<int64_t>(1024);
+    ++unit;
+  }
+
+  std::string result = absl::StrFormat("%3f", num_bytes / 1024.0);
+  std::string_view result_view = result;
+  while (absl::EndsWith(result_view, "0")) result_view.remove_suffix(1);
+  if (absl::EndsWith(result_view, ".")) result_view.remove_suffix(1);
+  return absl::StrFormat("%s%ciB", result_view, *unit);
+}
+
+/// Formats a list entry for printing according to the brief/human_readable
+/// flags.
+struct Formatter {
+  bool brief = false;
+  bool human_readable = false;
+
+  void SetBrief() {
+    this->brief = true;
+    this->human_readable = false;
+  }
+  void SetHumanReadable() {
+    this->human_readable = true;
+    this->brief = false;
+  }
+
+  std::string operator()(size_t width, const kvstore::ListEntry& entry) const {
+    if (brief || entry.size < 0) {
+      return absl::StrFormat("%s", entry.key);
+    }
+    if (!human_readable) {
+      return absl::StrFormat("%-*s%d", width + 2, entry.key, entry.size);
+    }
+    return absl::StrFormat("%-*s%s", width + 2, entry.key,
+                           FormatBytesHumanReadable(entry.size));
+  }
+};
+
 }  // namespace
-absl::Status KvstoreList(Context context,
-                         tensorstore::kvstore::Spec source_spec,
-                         tensorstore::span<std::string_view> args) {
+
+absl::Status KvstoreList(
+    Context context, tensorstore::kvstore::Spec source_spec,
+    tensorstore::span<std::string_view> args,
+    absl::FunctionRef<std::string(size_t, const kvstore::ListEntry&)>
+        formatter) {
   TENSORSTORE_ASSIGN_OR_RETURN(auto source,
                                kvstore::Open(source_spec, context).result());
 
@@ -88,18 +147,25 @@ absl::Status KvstoreList(Context context,
   TENSORSTORE_ASSIGN_OR_RETURN(auto list_entries,
                                kvstore::ListFuture(source, options).result());
 
+  size_t max_width = 0;
   RE2 re2(re_string);
   for (const auto& entry : list_entries) {
     if (re_string.empty() || RE2::FullMatch(entry.key, re2)) {
-      std::cout << entry.key << std::endl;
+      max_width = std::max(max_width, entry.key.size());
+    }
+  }
+  for (const auto& entry : list_entries) {
+    if (re_string.empty() || RE2::FullMatch(entry.key, re2)) {
+      std::cout << formatter(max_width, entry) << std::endl;
     }
   }
+
   return absl::OkStatus();
 }
 
 absl::Status RunKvstoreList(Context::Spec context_spec, CommandFlags flags) {
   tensorstore::JsonAbslFlag<std::optional<tensorstore::kvstore::Spec>> source;
-  std::vector<LongOption> options({
+  std::vector<LongOption> long_options({
       LongOption{"--source",
                  [&](std::string_view value) {
                    std::string error;
@@ -110,13 +176,23 @@ absl::Status RunKvstoreList(Context::Spec context_spec, CommandFlags flags) {
                  }},
   });
 
-  TENSORSTORE_RETURN_IF_ERROR(TryParseOptions(flags, options, {}));
+  Formatter formatter;
+  std::vector<BoolOption> bool_options({
+      BoolOption{"-h", [&]() { formatter.SetHumanReadable(); }},
+      BoolOption{"--human", [&]() { formatter.SetHumanReadable(); }},
+      BoolOption{"-b", [&]() { formatter.SetBrief(); }},
+      BoolOption{"--brief", [&]() { formatter.SetBrief(); }},
+  });
+
+  TENSORSTORE_RETURN_IF_ERROR(
+      TryParseOptions(flags, long_options, bool_options));
 
   tensorstore::Context context(context_spec);
 
   if (source.value) {
     if (flags.positional_args.empty()) flags.positional_args.push_back("");
-    return KvstoreList(context, *source.value, flags.positional_args);
+    return KvstoreList(context, *source.value, flags.positional_args,
+                       formatter);
   }
   if (flags.positional_args.empty()) {
     return absl::InvalidArgumentError(
@@ -128,13 +204,13 @@ absl::Status RunKvstoreList(Context::Spec context_spec, CommandFlags flags) {
   for (const auto spec : flags.positional_args) {
     auto from_url = kvstore::Spec::FromUrl(spec);
     if (from_url.ok()) {
-      status.Update(KvstoreList(context, from_url.value(), args));
+      status.Update(KvstoreList(context, from_url.value(), args, formatter));
       continue;
     }
     tensorstore::JsonAbslFlag<tensorstore::kvstore::Spec> arg_spec;
     std::string error;
     if (AbslParseFlag(spec, &arg_spec, &error)) {
-      status.Update(KvstoreList(context, arg_spec.value, args));
+      status.Update(KvstoreList(context, arg_spec.value, args, formatter));
       continue;
     }
     std::cerr << "Invalid spec: " << spec << ": " << error << std::endl;

tensorstore/tscli/main.cc

@@ -56,9 +56,9 @@ const char kUsage[] =
     "      Copies keys from one kvstore to another.\n"
     "  copy --source <source-kvstore-spec> --target <target-kvstore-spec>\n"
     "\n"
-    "  list"
-    "      List keys in a kvstore."
-    "  list --source <source-kvstore-spec> [glob...] \n"
+    "  list [-h/--human][-b/--brief]\n"
+    "      List keys in a kvstore.\n"
+    "  list --source <source-kvstore-spec> [glob...]\n"
     "  list [spec...]\n"
     "\n"
     "  search [-b/--brief][-f/--full]\n"