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)