Correctly set _paths for implementation of __repr__().

PiperOrigin-RevId: 516271701
google · Mar 13, 2023 · 9474099 · 9474099
1 parent 7ede6e7
commit 9474099
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/python/array_record_data_source.py b/python/array_record_data_source.py
@@ -217,6 +217,7 @@ def __init__(
           "a Sequence, String, pathlib.Path or FileInstruction."
       )
     self._read_instructions = _get_read_instructions(paths)
+    self._paths = [ri.filename for ri in self._read_instructions]
     # We open readers lazily when we need to read from them.
     self._readers = [None] * len(self._read_instructions)
     self._num_records = sum(

diff --git a/python/array_record_data_source_test.py b/python/array_record_data_source_test.py
@@ -174,6 +174,13 @@ def test_array_record_source_empty_sequence(self):
     with self.assertRaises(ValueError):
       array_record_data_source.ArrayRecordDataSource([])
 
+  def test_repr(self):
+    ar = array_record_data_source.ArrayRecordDataSource([
+        self.testdata_dir / "digits.array_record-00000-of-00002",
+        self.testdata_dir / "digits.array_record-00001-of-00002",
+    ])
+    self.assertRegex(repr(ar), r"ArrayRecordDataSource\(hash_of_paths=[\w]+\)")
+
 
 class RunInParallelTest(parameterized.TestCase):