Merge pull request tensorflow#692 from us:document_dataset_bug

PiperOrigin-RevId: 257705157

tensorflow_datasets/core/features/features_dict.py

@@ -138,7 +138,13 @@ def __iter__(self):
 
   def __repr__(self):
     """Display the feature dictionary."""
-    return '{}({})'.format(type(self).__name__, self._feature_dict)
+    lines = ['{}({{'.format(type(self).__name__)]
+    # Add indentation
+    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."""

tensorflow_datasets/core/features/features_test.py

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

tensorflow_datasets/core/features/sequence_feature.py

@@ -194,9 +194,13 @@ def __setstate__(self, state):
     del state['__getattr__']
     self.__dict__.update(state)
 
-  def _additional_repr_info(self):
-    """Override to return additional info to go into __repr__."""
-    return {'feature': repr(self._feature)}
+  def __repr__(self):
+    """Display the feature."""
+    inner_feature_repr = repr(self._feature)
+    if inner_feature_repr.startswith('FeaturesDict('):
+      # Minor formatting cleaning: 'Sequence(FeaturesDict({' => 'Sequence({'
+      inner_feature_repr = inner_feature_repr[len('FeaturesDict('):-len(')')]
+    return '{}({})'.format(type(self).__name__, inner_feature_repr)
 
 
 def stack_arrays(*elems):

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):
-  """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):
-      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):