no-op refactorting: file_format_adapter write_from_generator accepts a generator, instead of generator_fn.

PiperOrigin-RevId: 251828067

Author: pierrot0

tensorflow_datasets/core/dataset_builder.py

@@ -823,8 +823,7 @@ def _prepare_split(self, split_generator, max_examples_per_split):
       generator = itertools.islice(generator, max_examples_per_split)
     generator = (self.info.features.encode_example(ex) for ex in generator)
     output_files = self._build_split_filenames(split_generator.split_info)
-    self._file_format_adapter.write_from_generator(
-        lambda: generator, output_files)
+    self._file_format_adapter.write_from_generator(generator, output_files)
 
 
 class BeamBasedBuilder(FileAdapterBuilder):

tensorflow_datasets/core/file_format_adapter.py

@@ -69,12 +69,11 @@ def __init__(self, example_specs):
     del example_specs
 
   @abc.abstractmethod
-  def write_from_generator(self, generator_fn, output_files):
+  def write_from_generator(self, generator, output_files):
     """Write to files from generators_and_filenames.
 
     Args:
-      generator_fn: returns generator yielding dictionaries of feature name to
-        value.
+      generator: generator yielding dictionaries of feature name to value.
       output_files: `list<str>`, output files to write files to.
     """
     raise NotImplementedError
@@ -121,10 +120,9 @@ def __init__(self, example_specs):
         example_specs)
     self._parser = example_parser.ExampleParser(example_specs)
 
-  def write_from_generator(self, generator_fn, output_files):
-    wrapped = (
-        self._serializer.serialize_example(example)
-        for example in generator_fn())
+  def write_from_generator(self, generator, output_files):
+    wrapped = (self._serializer.serialize_example(example)
+               for example in generator)
     _write_tfrecords_from_generator(wrapped, output_files, shuffle=True)
 
   def write_from_pcollection(self, pcollection, file_path_prefix, num_shards):

tensorflow_datasets/testing/test_utils.py

@@ -296,7 +296,7 @@ def features_encode_decode(features_dict, example, as_tensor=False):
     file_adapter = file_format_adapter.TFRecordExampleAdapter(
         features_dict.get_serialized_info())
     file_adapter.write_from_generator(
-        generator_fn=lambda: [encoded_example],
+        generator=[encoded_example],
         output_files=[tmp_filename],
     )
     ds = file_adapter.dataset_from_filename(tmp_filename)