Allow underscore in slice value syntax (eg 50_000)

PiperOrigin-RevId: 403945855
sputney13 · Oct 18, 2021 · 0f97767 · 0f97767
1 parent bf78f7e
commit 0f97767
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/tensorflow_datasets/core/tfrecords_reader.py b/tensorflow_datasets/core/tfrecords_reader.py
@@ -64,7 +64,7 @@
 _SLICE_RE = re.compile(
     r"""^
     (
-        (?P<val>-?\d+)
+        (?P<val>-?[\d_]+)
         (?P<unit>(?:%|shard))?
     )?
     $""",

diff --git a/tensorflow_datasets/core/tfrecords_reader_test.py b/tensorflow_datasets/core/tfrecords_reader_test.py
@@ -249,6 +249,9 @@ def test_valid(self):
     # Empty:
     # Slices resulting in empty datasets are valid with 'closest' rounding:
     self.check_from_spec('validation[:1%]', [('validation', None, 0)])
+    # New integer syntax
+    self.check_from_spec('train[4_2:9_9]', [('train', 42, 99)])
+    self.check_from_spec('train[:1_0%]', [('train', None, 20)])
 
     # Supports splits with '-' in name.
     ri = self.check_from_spec('dev-train', [('dev-train', None, None)])