Change FeaturesDict.__repr__ method.

Move `pprint_features_dict` func to `features_dict.py`

Author: us

tensorflow_datasets/core/features/features_dict.py

@@ -22,11 +22,11 @@
 
 import six
 import tensorflow as tf
+import tensorflow_datasets as tfds
 
 from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.features import feature as feature_lib
 from tensorflow_datasets.core.features import top_level_feature
-from tensorflow_datasets.scripts.document_datasets import pprint_features_dict
 
 
 class FeaturesDict(top_level_feature.TopLevelFeature):
@@ -139,7 +139,13 @@ def __iter__(self):
 
   def __repr__(self):
     """Display the feature dictionary."""
-    return pprint_features_dict(self._feature_dict, types='FeaturesDict')
+    lines = ['{}({{'.format(type(self).__name__)]
+    for key, feature in sorted(list(self._feature_dict.items())):
+      all_sub_lines = '\'{}\': {},'.format(key, feature)
+      lines.extend('    ' + l for l in
+                   all_sub_lines.split('\n'))
+    lines.append('})')
+    return '\n'.join(lines)
 
   def get_tensor_info(self):
     """See base class for details."""
@@ -232,3 +238,26 @@ def to_feature(value):
     return FeaturesDict(value)
   else:
     raise ValueError('Feature not supported: {}'.format(value))
+
+
+def pprint_features_dict(features_dict, indent=0, add_prefix=True):
+  """Pretty-print tfds.features.FeaturesDict."""
+  first_last_indent_str = " " * indent
+  indent_str = " " * (indent + 4)
+  first_line = "%s%s({" % (
+      first_last_indent_str if add_prefix else "",
+      type(features_dict).__name__,
+  )
+  lines = [first_line]
+  for k in sorted(list(features_dict.keys())):
+    v = features_dict[k]
+    if isinstance(v, tfds.features.FeaturesDict) or (
+        isinstance(v, tfds.features.Sequence) and
+        isinstance(v, tfds.features.FeaturesDict)
+    ):
+      v_str = pprint_features_dict(v, indent + 4, False)
+    else:
+      v_str = str(v)
+    lines.append("%s'%s': %s," % (indent_str, k, v_str))
+  lines.append("%s})" % first_last_indent_str)
+  return "\n".join(lines)

tensorflow_datasets/core/features/features_test.py

@@ -196,6 +196,25 @@ def test_feature_getitem(self):
     self.assertEqual(fdict['integer'].dtype, tf.int32)
     self.assertEqual(fdict['string'].dtype, tf.string)
 
+  def test_feature__repr__(self):
+
+    FEATURE_STR = (\
+      "FeaturesDict({\n"
+      "    'label': Sequence(ClassLabel(shape=(), dtype=tf.int64, num_classes=2)),\n"
+      "    'metadata': Sequence(FeaturesDict({\n"
+      "        'frame': Image(shape=(32, 32, 3), dtype=tf.uint8),\n"
+      "    })),\n"
+      "})")
+
+    feature_dict = features_lib.FeaturesDict({
+      'metadata': features_lib.Sequence({
+        'frame': features_lib.Image(shape=(32, 32, 3)),
+      }),
+      'label': features_lib.Sequence(features_lib.ClassLabel(names=["m", "f"])),
+    })
+
+    self.assertEqual(str(feature_dict), FEATURE_STR)
+
 
 class FeatureTensorTest(testing.FeatureExpectationsTestCase):
 

tensorflow_datasets/scripts/document_datasets.py

@@ -311,29 +311,9 @@ def make_module_to_builder_dict(datasets=None):
   return module_to_builder
 
 
-def pprint_features_dict(features_dict, indent=0, add_prefix=True, types=None):
-  """Pretty-print tfds.features.FeaturesDict."""
-  first_last_indent_str = " " * indent
-  indent_str = " " * (indent + 4)
-  first_line = "%s%s({" % (
-      first_last_indent_str if add_prefix else "",
-      type(features_dict).__name__ if type is None else types,
-  )
-  lines = [first_line]
-  for k in sorted(list(features_dict.keys())):
-    v = features_dict[k]
-    if hasattr(v, 'keys'):
-      v_str = pprint_features_dict(v, indent + 4, False)
-    else:
-      v_str = str(v)
-    lines.append("%s'%s': %s," % (indent_str, k, v_str))
-  lines.append("%s})" % first_last_indent_str)
-  return "\n".join(lines)
-
-
 def make_feature_information(info):
   """Make feature information table."""
-  return FEATURE_BLOCK % pprint_features_dict(info.features)
+  return FEATURE_BLOCK % info.features
 
 
 def make_citation(citation):