Minor DatasetInfo print formatting change

PiperOrigin-RevId: 258087427

Author: Conchylicultor

tensorflow_datasets/core/dataset_info.py

@@ -39,7 +39,6 @@
 import json
 import os
 import posixpath
-import pprint
 import tempfile
 
 from absl import logging
@@ -72,7 +71,7 @@
     total_num_examples={total_num_examples},
     splits={splits},
     supervised_keys={supervised_keys},
-    citation='{citation}',
+    citation={citation},
     redistribution_info={redistribution_info},
 )
 """
@@ -408,19 +407,13 @@ def initialize_from_bucket(self):
       gcs_utils.download_gcs_file(fname, out_fname)
     self.read_from_directory(tmp_dir)
 
-  def __str__(self):
-    splits_pprint = "{\n %s\n    }" % (
-        pprint.pformat(
-            {k: self.splits[k] for k in sorted(list(self.splits.keys()))},
-            indent=8, width=1)[1:-1])
-    features_dict = self.features
-    features_pprint = "%s({\n %s\n    }" % (
-        type(features_dict).__name__,
-        pprint.pformat({
-            k: features_dict[k] for k in sorted(list(features_dict.keys()))
-        }, indent=8, width=1)[1:-1])
-    citation_pprint = '"""\n%s\n    """' % "\n".join(
-        [u" " * 8 + line for line in self.citation.split(u"\n")])
+  def __repr__(self):
+    splits_pprint = _indent("\n".join(["{"] + [
+        "    '{}': {},".format(k, split.num_examples)
+        for k, split in sorted(self.splits.items())
+    ] + ["}"]))
+    features_pprint = _indent(repr(self.features))
+    citation_pprint = _indent('"""{}"""'.format(self.citation.strip()))
     return INFO_STR.format(
         name=self.name,
         version=self.version,
@@ -431,7 +424,14 @@ def __str__(self):
         citation=citation_pprint,
         urls=self.urls,
         supervised_keys=self.supervised_keys,
-        redistribution_info=self.redistribution_info)
+        # Proto add a \n that we strip.
+        redistribution_info=str(self.redistribution_info).strip())
+
+
+def _indent(content):
+  """Add indentation to all lines except the first."""
+  lines = content.split("\n")
+  return "\n".join([lines[0]] + ["    " + l for l in lines[1:]])
 
 #
 #

tensorflow_datasets/core/dataset_info_test.py

@@ -23,6 +23,7 @@
 import os
 import tempfile
 import numpy as np
+import six
 import tensorflow as tf
 from tensorflow_datasets import testing
 from tensorflow_datasets.core import dataset_info
@@ -76,6 +77,7 @@ def setUpClass(cls):
 
   @classmethod
   def tearDownClass(cls):
+    super(DatasetInfoTest, cls).tearDownClass()
     testing.rm_tmp_dir(cls._tfds_tmp_dir)
 
   def test_undefined_dir(self):
@@ -106,8 +108,8 @@ def test_reading(self):
     self.assertTrue(len(split_dict), 2)
 
     # Assert on what they are
-    self.assertTrue("train" in split_dict)
-    self.assertTrue("test" in split_dict)
+    self.assertIn("train", split_dict)
+    self.assertIn("test", split_dict)
 
     # Assert that this is computed correctly.
     self.assertEqual(40, info.splits.total_num_examples)
@@ -127,7 +129,8 @@ def test_writing(self):
     mnist_builder = mnist.MNIST(
         data_dir=tempfile.mkdtemp(dir=self.get_temp_dir()))
 
-    info = dataset_info.DatasetInfo(builder=mnist_builder)
+    info = dataset_info.DatasetInfo(
+        builder=mnist_builder, features=mnist_builder.info.features)
     info.read_from_directory(_INFO_DIR)
 
     # Read the json file into a string.
@@ -152,6 +155,10 @@ def test_writing(self):
     # Assert correct license was written.
     self.assertEqual(existing_json["redistributionInfo"]["license"], license_)
 
+    if six.PY3:
+      # Only test on Python 3 to avoid u'' formatting issues
+      self.assertEqual(repr(info), INFO_STR)
+
   def test_restore_after_modification(self):
     # Create a DatasetInfo
     info = dataset_info.DatasetInfo(
@@ -296,5 +303,32 @@ def test_updates_on_bucket_info(self):
     self.assertEqual(2, len(info.as_proto.schema.feature))
 
 
+INFO_STR = """tfds.core.DatasetInfo(
+    name='mnist',
+    version=1.0.0,
+    description='The MNIST database of handwritten digits.',
+    urls=['https://storage.googleapis.com/cvdf-datasets/mnist/'],
+    features=FeaturesDict({
+        'image': Image(shape=(28, 28, 1), dtype=tf.uint8),
+        'label': ClassLabel(shape=(), dtype=tf.int64, num_classes=10),
+    }),
+    total_num_examples=40,
+    splits={
+        'test': 20,
+        'train': 20,
+    },
+    supervised_keys=('image', 'label'),
+    citation=\"\"\"@article{lecun2010mnist,
+      title={MNIST handwritten digit database},
+      author={LeCun, Yann and Cortes, Corinna and Burges, CJ},
+      journal={ATT Labs [Online]. Available: http://yann. lecun. com/exdb/mnist},
+      volume={2},
+      year={2010}
+    }\"\"\",
+    redistribution_info=license: "test license",
+)
+"""
+
+
 if __name__ == "__main__":
   testing.test_main()

tensorflow_datasets/core/splits.py

@@ -36,7 +36,7 @@ class SplitInfo(object):
 
   @property
   def num_examples(self):
-    return self.statistics.num_examples
+    return int(self.statistics.num_examples)
 
   def __repr__(self):
     num_examples = self.num_examples or "unknown"