def _roundtrip(tmpdir,
tree,
compression=None,
write_options={},
read_options={}):
tmpfile = os.path.join(str(tmpdir), 'test.asdf')
ff = asdf.AsdfFile(tree)
ff.set_array_compression(tree['science_data'], compression)
ff.write_to(tmpfile, **write_options)
with asdf.AsdfFile.open(tmpfile, mode="rw") as ff:
ff.update(**write_options)
with asdf.AsdfFile.open(tmpfile, **read_options) as ff:
helpers.assert_tree_match(tree, ff.tree)
# Also test saving to a buffer
buff = io.BytesIO()
ff = asdf.AsdfFile(tree)
ff.set_array_compression(tree['science_data'], compression)
ff.write_to(buff, **write_options)
buff.seek(0)
with asdf.AsdfFile.open(buff, **read_options) as ff:
helpers.assert_tree_match(tree, ff.tree)
# Test saving to a non-seekable buffer
buff = io.BytesIO()
ff = asdf.AsdfFile(tree)
ff.set_array_compression(tree['science_data'], compression)
ff.write_to(generic_io.OutputStream(buff), **write_options)
buff.seek(0)
with asdf.AsdfFile.open(generic_io.InputStream(buff),
**read_options) as ff:
helpers.assert_tree_match(tree, ff.tree)
return ff
def test_stream_to_stream():
tree = {
'nonstream': np.array([1, 2, 3, 4], np.int64),
'stream': stream.Stream([6, 2], np.float64)
}
buff = io.BytesIO()
fd = generic_io.OutputStream(buff)
ff = asdf.AsdfFile(tree)
ff.write_to(fd)
for i in range(100):
fd.write(np.array([i] * 12, np.float64).tostring())
buff.seek(0)
with asdf.AsdfFile().open(generic_io.InputStream(buff, 'r')) as ff:
assert len(ff.blocks) == 2
assert_array_equal(ff.tree['nonstream'], np.array([1, 2, 3, 4], np.int64))
assert ff.tree['stream'].shape == (100, 6, 2)
for i, row in enumerate(ff.tree['stream']):
assert np.all(row == i)
def get_write_fd():
return generic_io.OutputStream(buff)
