Add special format for FeaturesDict.__repr__

`tensorflow_datasets.scripts.document_datasets.pprint_features_dict`
converted to public function and add types parameter.

Author: us

tensorflow_datasets/core/features/features_dict.py

@@ -26,6 +26,7 @@
 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):
@@ -138,7 +139,7 @@ def __iter__(self):
 
   def __repr__(self):
     """Display the feature dictionary."""
-    return '{}({})'.format(type(self).__name__, self._feature_dict)
+    return pprint_features_dict(self._feature_dict, types='FeaturesDict')
 
   def get_tensor_info(self):
     """See base class for details."""

tensorflow_datasets/scripts/document_datasets.py

@@ -311,19 +311,19 @@ def make_module_to_builder_dict(datasets=None):
   return module_to_builder
 
 
-def _pprint_features_dict(features_dict, indent=0, add_prefix=True):
+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__,
+      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)
+      v_str = pprint_features_dict(v, indent + 4, False)
     else:
       v_str = str(v)
     lines.append("%s'%s': %s," % (indent_str, k, v_str))
@@ -333,7 +333,7 @@ def _pprint_features_dict(features_dict, indent=0, add_prefix=True):
 
 def make_feature_information(info):
   """Make feature information table."""
-  return FEATURE_BLOCK % _pprint_features_dict(info.features)
+  return FEATURE_BLOCK % pprint_features_dict(info.features)
 
 
 def make_citation(citation):