1 files changed, 38 insertions, 0 deletions
diff --git a/startop/scripts/lib/print_utils.py b/startop/scripts/lib/print_utils.py
index c33e0f979a0c..8c5999d99d6e 100644
--- a/startop/scripts/lib/print_utils.py
+++ b/startop/scripts/lib/print_utils.py
@@ -27,3 +27,41 @@ def debug_print(*args, **kwargs):
 
 def error_print(*args, **kwargs):
   print('[ERROR]:', *args, file=sys.stderr, **kwargs)
+
+def _expand_gen_repr(args):
+  """Like repr but any generator-like object has its iterator consumed
+  and then called repr on."""
+  new_args_list = []
+  for i in args:
+    # detect iterable objects that do not have their own override of __str__
+    if hasattr(i, '__iter__'):
+      to_str = getattr(i, '__str__')
+      if to_str.__objclass__ == object:
+        # the repr for a generator is just type+address, expand it out instead.
+        new_args_list.append([_expand_gen_repr([j])[0] for j in i])
+        continue
+    # normal case: uses the built-in to-string
+    new_args_list.append(i)
+  return new_args_list
+
+def debug_print_gen(*args, **kwargs):
+  """Like _debug_print but will turn any iterable args into a list."""
+  if not DEBUG:
+    return
+
+  new_args_list = _expand_gen_repr(args)
+  debug_print(*new_args_list, **kwargs)
+
+def debug_print_nd(*args, **kwargs):
+  """Like _debug_print but will turn any NamedTuple-type args into a string."""
+  if not DEBUG:
+    return
+
+  new_args_list = []
+  for i in args:
+    if hasattr(i, '_field_types'):
+      new_args_list.append("%s: %s" % (i.__name__, i._field_types))
+    else:
+      new_args_list.append(i)
+
+  debug_print(*new_args_list, **kwargs)