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)