def default_writers(output_path): # pylint: disable=invalid-name """Returns the default writers for use in WriteResults. Args: output_path: Path to store results files under. """ writers = [] output = { constants.METRICS_KEY: os.path.join(output_path, _METRICS_OUTPUT_FILE), constants.PLOTS_KEY: os.path.join(output_path, _PLOTS_OUTPUT_FILE) } for (key, output_file) in output.items(): writers.append( writer.Writer(stage_name='WriteTFRecord(%s)' % output_file, ptransform=writer.Write( key=key, ptransform=beam.io.WriteToTFRecord( file_path_prefix=output_file, shard_name_template='')))) return writers
def testWriteIgnoresMissingKeys(self): with beam.Pipeline() as pipeline: test = pipeline | beam.Create(['test']) # PTransform is None so this will throw exception if it tries to run _ = {'test': test} | writer.Write('key-does-not-exist', None)