def configure_output_stream(dburls, mapping_file, output_format, output_files,
output_folder):
assert isinstance(output_files, (list, type(None)))
output_streams = [] # we allow multiple output streams
for dburl in dburls:
if mapping_file:
with click.open_file(mapping_file, "r") as f:
mappings = yaml.safe_load(f)
else:
mappings = None
output_streams.append(SqlOutputStream.from_url(dburl, mappings))
# JSON is the only output format (other than debug) that can go on stdout
if output_format == "json" and not output_files:
output_streams.append(JSONOutputStream(sys.stdout))
if output_format == "csv":
output_streams.append(CSVOutputStream(output_folder))
if output_files:
for path in output_files:
if output_folder:
path = Path(output_folder,
path) # put the file in the output folder
format = output_format or Path(path).suffix[1:]
if format == "json":
output_streams.append(JSONOutputStream(path))
elif format == "txt":
output_streams.append(DebugOutputStream(path))
elif format == "dot":
output_streams.append(GraphvizOutputStream(path))
elif format in graphic_file_extensions:
output_streams.append(ImageOutputStream(path, format))
else:
raise click.ClickException(
f"Unknown format or file extension: {format}")
if len(output_streams) == 0:
output_stream = DebugOutputStream()
elif len(output_streams) == 1:
output_stream = output_streams[0]
else:
output_stream = MultiplexOutputStream(output_streams)
try:
yield output_stream
finally:
for output_stream in output_streams:
try:
messages = output_stream.close()
except Exception as e:
messages = None
click.echo(f"Could not close {output_stream}: {str(e)}",
err=True)
if messages:
for message in messages:
click.echo(message)
def standard_runtime():
output_stream = DebugOutputStream()
interpreter = Interpreter(
output_stream=output_stream,
parent_application=SnowfakeryApplication(),
parse_result=FakeParseResult(),
globals=Globals(),
)
runtime_context = RuntimeContext(interpreter=interpreter)
interpreter.current_context = runtime_context
return runtime_context
def test_fail_render_weird_type(self):
with self.assertRaises((DataGenError, TypeError)):
o = ObjectTemplate(
"abcd",
filename="abc.yml",
line_num=10,
fields=[
FieldFactory(
"x",
SimpleValue(b"junk", filename="abc.yml", line_num=42),
**line)
],
)
o.generate_rows(DebugOutputStream(), standard_runtime())
def test_fail_render_weird_template(self):
with pytest.raises(DataGenError):
o = ObjectTemplate(
"abcd",
filename="abc.yml",
line_num=10,
fields=[
FieldFactory(
"x",
SimpleValue("${{5()}}", filename="abc.yml", line_num=42),
**line
)
],
)
o.generate_rows(DebugOutputStream(), standard_runtime())
def configure_output_stream(dburls, output_format, output_files, output_folder,
parent_application):
assert isinstance(output_files, (list, type(None)))
with _get_output_streams(dburls, output_files, output_format,
output_folder) as output_streams:
if len(output_streams) == 0:
output_stream = DebugOutputStream()
elif len(output_streams) == 1:
output_stream = output_streams[0]
else:
output_stream = MultiplexOutputStream(output_streams)
try:
yield output_stream
finally:
try:
messages = output_stream.close()
except Exception as e:
messages = None
parent_application.echo(
f"Could not close {output_stream}: {str(e)}", err=True)
if messages:
for message in messages:
parent_application.echo(message)
def test_render_empty_object_template(self):
o = ObjectTemplate("abcd", filename="abc.yml", line_num=10)
o.generate_rows(DebugOutputStream(), standard_runtime())
def standard_runtime():
output_stream = DebugOutputStream()
return RuntimeContext(Interpreter(output_stream=output_stream), "Foo")
def standard_runtime():
output_stream = DebugOutputStream()
interpreter = Interpreter(output_stream=output_stream, globals=Globals())
runtime_context = RuntimeContext(interpreter=interpreter)
interpreter.current_context = runtime_context
return runtime_context