def _write_tree(self, tree: dict, output: Union[str, BinaryIO], file_mode: int = 0o666) -> None: """ Write the model to disk. :param tree: The data dict - will be the ASDF tree. :param output: The output file path or a file object. :param file_mode: The output file's permissions. :return: None """ self.meta["created_at"] = get_datetime_now() meta = self.meta.copy() meta["environment"] = collect_environment() final_tree = {} final_tree.update(tree) final_tree["meta"] = meta isfileobj = not isinstance(output, str) if not isfileobj: self._source = output path = output output = open(output, "wb") os.chmod(path, file_mode) pos = 0 else: pos = output.tell() try: with asdf.AsdfFile(final_tree) as file: queue = [("", tree)] while queue: path, element = queue.pop() if isinstance(element, dict): for key, val in element.items(): queue.append((path + "/" + key, val)) elif isinstance(element, (list, tuple)): for child in element: queue.append((path, child)) elif isinstance(element, numpy.ndarray): path += "/" if path not in self._compression_prefixes: self._log.debug("%s -> %s compression", path, self.ARRAY_COMPRESSION) file.set_array_compression(element, self.ARRAY_COMPRESSION) else: self._log.debug("%s -> compression disabled", path) file.write_to(output) self._size = output.seek(0, os.SEEK_END) - pos finally: if not isfileobj: output.close()
def test_stream_real_file(tmpdir): path = os.path.join(str(tmpdir), 'test.asdf') tree = { 'nonstream': np.array([1, 2, 3, 4], np.int64), 'stream': stream.Stream([6, 2], np.float64) } with open(path, 'wb') as fd: ff = asdf.AsdfFile(tree) ff.write_to(fd) for i in range(100): fd.write(np.array([i] * 12, np.float64).tostring()) with asdf.AsdfFile().open(path) as ff: assert len(ff.blocks) == 1 assert_array_equal(ff.tree['nonstream'], np.array([1, 2, 3, 4], np.int64)) assert ff.tree['stream'].shape == (100, 6, 2) assert len(ff.blocks) == 2 for i, row in enumerate(ff.tree['stream']): assert np.all(row == i)
def test_arbitrary_python_object(): # Putting "just any old" Python object in the tree should raise an # exception. class Foo(object): pass tree = {'object': Foo()} buff = io.BytesIO() ff = asdf.AsdfFile(tree) with pytest.raises(yaml.YAMLError): ff.write_to(buff)
def output_to_asdf(output, input, label='output', outdir='.'): fname_prefix = input.get('prefix', '') fname_suffix = input.get('suffix', '') with asdf.AsdfFile() as f: f.tree.update(output) f.tree['input'] = copy.deepcopy(input) f.tree['experiment'] = copy.deepcopy(input['experiment']) fname = f'{fname_prefix}{label}{fname_suffix}.asdf' os.makedirs(outdir, exist_ok=True) # pdb.set_trace() f.write_to(outdir + "/" + fname) print(fname)
def test_unicode_to_list(tmpdir): arr = np.array(['', '𐀠'], dtype='<U') tree = {'unicode': arr} fd = io.BytesIO() ff = asdf.AsdfFile(tree) ff.set_array_storage(arr, 'inline') ff.write_to(fd) fd.seek(0) with asdf.open(fd) as ff: ff.resolve_and_inline() ff.write_to(io.BytesIO())
def test_checksum(tmpdir): tmpdir = str(tmpdir) path = os.path.join(tmpdir, 'test.asdf') my_array = np.arange(0, 64, dtype='<i8').reshape((8, 8)) tree = {'my_array': my_array} ff = asdf.AsdfFile(tree) ff.write_to(path, auto_inline=None) with asdf.open(path, validate_checksums=True) as ff: assert type(ff.blocks._internal_blocks[0].checksum) == bytes assert ff.blocks._internal_blocks[0].checksum == \ b'\xcaM\\\xb8t_L|\x00\n+\x01\xf1\xcfP1'
def test_numpy_scalar(numpy_value, expected_value): ctx = asdf.AsdfFile() tree = {"value": numpy_value} buffer = io.BytesIO() yamlutil.dump_tree(tree, buffer, ctx) buffer.seek(0) if isinstance(expected_value, float) and NUMPY_LT_1_14: assert yamlutil.load_tree(buffer)["value"] == pytest.approx( expected_value, rel=0.001) else: assert yamlutil.load_tree(buffer)["value"] == expected_value
def test_make_reference(tmpdir): exttree = { # Include some ~ and / in the name to make sure that escaping # is working correctly 'f~o~o/': { 'a': np.array([0, 1, 2], np.float), 'b': np.array([3, 4, 5], np.float) } } external_path = os.path.join(str(tmpdir), 'external.asdf') ext = asdf.AsdfFile(exttree) ext.write_to(external_path) with asdf.open(external_path) as ext: ff = asdf.AsdfFile() ff.tree['ref'] = ext.make_reference(['f~o~o/', 'a']) assert_array_equal(ff.tree['ref'], ext.tree['f~o~o/']['a']) ff.write_to(os.path.join(str(tmpdir), 'source.asdf')) with asdf.open(os.path.join(str(tmpdir), 'source.asdf')) as ff: assert ff.tree['ref']._uri == 'external.asdf#f~0o~0o~1/a'
def test_custom_validation_with_definitions_good(tmpdir): custom_schema_path = helpers.get_test_data_path( 'custom_schema_definitions.yaml') asdf_file = str(tmpdir.join('out.asdf')) # This tree conforms to the custom schema tree = {'thing': {'biz': 'hello', 'baz': 'world'}} with asdf.AsdfFile(tree, custom_schema=custom_schema_path) as ff: ff.write_to(asdf_file) with asdf.open(asdf_file, custom_schema=custom_schema_path): pass
def test_max_min_literals(num): tree = { 'test_int': num, } with pytest.raises(ValidationError): asdf.AsdfFile(tree) tree = { 'test_list': [num], } with pytest.raises(ValidationError): asdf.AsdfFile(tree) tree = { num: 'test_key', } with pytest.raises(ValidationError): asdf.AsdfFile(tree)
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 to_asdf(self, filename): """ Serialize an `EmissionModel` to an ASDF file """ tree = { 'temperature': self.temperature, 'density': self.density, 'ions': [ion.ion_name for ion in self], 'dset_names': [ion._dset_names for ion in self] } tree['emissivity_table_filename'] = self.emissivity_table_filename with asdf.AsdfFile(tree) as asdf_file: asdf_file.write_to(filename)
def assert_model_roundtrip(model, tmpdir, version=None): """ Assert that a model can be written to an ASDF file and read back in without losing any of its essential properties. """ path = str(tmpdir / "test.asdf") with asdf.AsdfFile({"model": model}, version=version) as af: af.write_to(path) with asdf.open(path) as af: assert_models_equal(model, af["model"]) return af["model"]
def test_no_blocks(tmp_path, create_editor, version): file_path = str(tmp_path / "test.asdf") with asdf.AsdfFile(version=version) as af: af["foo"] = "bar" af.write_to(file_path) os.environ["EDITOR"] = create_editor(r"foo: bar", "foo: baz") assert main.main_from_args(["edit", file_path]) == 0 with asdf.open(file_path) as af: assert af["foo"] == "baz"
def test_copy_inline(): yaml = """ x0: !core/ndarray-1.0.0 data: [-1.0, 1.0] """ buff = helpers.yaml_to_asdf(yaml) with asdf.open(buff) as infile: with asdf.AsdfFile() as f: f.tree['a'] = infile.tree['x0'] f.tree['b'] = f.tree['a'] f.write_to(io.BytesIO())
def to_asdf(self, filename): """ Serialize this instance of `Skeleton` to an ASDF file """ tree = {} for l in self.loops: tree[l.name] = { 'field_strength': l.field_strength, 'coordinate': l.coordinate, 'model_results_filename': l.model_results_filename, } with asdf.AsdfFile(tree) as asdf_file: asdf_file.write_to(filename)
def test_is_asdf(tmpdir): # test fits fits = pytest.importorskip('astropy.io.fits') hdul = fits.HDUList() phdu = fits.PrimaryHDU() imhdu = fits.ImageHDU(data=np.arange(24).reshape((4, 6))) hdul.append(phdu) hdul.append(imhdu) path = os.path.join(str(tmpdir), 'test.fits') hdul.writeto(path) assert not is_asdf_file(path) assert is_asdf_file(asdf.AsdfFile())
def test_empty_split_save_load_merge(self): strings = [] merged = merge_strings(strings) assert_array_equal(merged["strings"], numpy.array([], dtype="S1")) assert_array_equal(merged["lengths"], numpy.array([], dtype=int)) self.assertIsNone(merged["str"]) af = asdf.AsdfFile(merged) buffer = BytesIO() af.write_to(buffer) buffer.seek(0) af_loaded = asdf.open(buffer) strings_restored = split_strings(af_loaded.tree) self.assertEqual(strings, strings_restored)
def test_no_block_index(): buff = io.BytesIO() arrays = [] for i in range(10): arrays.append(np.ones((8, 8)) * i) tree = {'arrays': arrays} ff = asdf.AsdfFile(tree) ff.write_to(buff, include_block_index=False) assert constants.INDEX_HEADER not in buff.getvalue()
def test_recompression(tmpdir): tree = _get_large_tree() tmpfile = os.path.join(str(tmpdir), 'test1.asdf') afile = asdf.AsdfFile(tree) afile.write_to(tmpfile, all_array_compression='zlib') afile.close() afile = asdf.AsdfFile.open(tmpfile) tmpfile = os.path.join(str(tmpdir), 'test2.asdf') afile.write_to(tmpfile, all_array_compression='bzp2') afile.close() afile = asdf.AsdfFile.open(tmpfile) helpers.assert_tree_match(tree, afile.tree) afile.close()
def test_update_expand_tree(tmpdir): tmpdir = str(tmpdir) testpath = os.path.join(tmpdir, "test.asdf") # This is the case where the new tree can't fit in the available space my_array = np.arange(64) * 1 my_array2 = np.arange(64) * 2 tree = {'arrays': [my_array, my_array2, np.arange(3)]} ff = asdf.AsdfFile(tree) ff.set_array_storage(tree['arrays'][2], 'inline') assert len(list(ff.blocks.inline_blocks)) == 1 ff.write_to(testpath, auto_inline=None, pad_blocks=True) with asdf.open(testpath, mode='rw') as ff: assert_array_equal(ff.tree['arrays'][0], my_array) orig_offset = ff.blocks[ff.tree['arrays'][0]].offset ff.tree['extra'] = [0] * 6000 ff.update() with asdf.open(testpath) as ff: assert orig_offset <= ff.blocks[ff.tree['arrays'][0]].offset assert ff.blocks[ff.tree['arrays'][2]].array_storage == 'inline' assert_array_equal(ff.tree['arrays'][0], my_array) assert_array_equal(ff.tree['arrays'][1], my_array2) # Now, we expand the header only by a little bit ff = asdf.AsdfFile(tree) ff.set_array_storage(tree['arrays'][2], 'inline') ff.write_to(os.path.join(tmpdir, "test2.asdf"), pad_blocks=True) with asdf.open(os.path.join(tmpdir, "test2.asdf"), mode='rw') as ff: orig_offset = ff.blocks[ff.tree['arrays'][0]].offset ff.tree['extra'] = [0] * 2 ff.update() with asdf.open(os.path.join(tmpdir, "test2.asdf")) as ff: assert orig_offset == ff.blocks[ff.tree['arrays'][0]].offset assert ff.blocks[ff.tree['arrays'][2]].array_storage == 'inline' assert_array_equal(ff.tree['arrays'][0], my_array) assert_array_equal(ff.tree['arrays'][1], my_array2)
def test_versioned_writing(): from ..tags.core.complex import ComplexType # Create a bogus version map versioning._version_map['42.0.0'] = { 'FILE_FORMAT': '42.0.0', 'YAML_VERSION': '1.1', 'tags': { 'tag:stsci.edu:asdf/core/complex': '42.0.0', 'tag:stscu.edu:asdf/core/asdf': '1.0.0' } } versioning.supported_versions.append(versioning.AsdfVersion('42.0.0')) class FancyComplexType(ComplexType, asdftypes.CustomType): version = (42, 0, 0) # This is a sanity check to ensure that the custom FancyComplexType does # not get added to ASDF's built-in extension, since this would cause any # subsequent tests that rely on ComplexType to fail. assert not issubclass(FancyComplexType, asdftypes.AsdfTypeMeta) class FancyComplexExtension(object): @property def types(self): return [FancyComplexType] @property def tag_mapping(self): return [] @property def url_mapping(self): return [ ('http://stsci.edu/schemas/asdf/core/complex-42.0.0', util.filepath_to_url(TEST_DATA_PATH) + '/complex-42.0.0.yaml') ] tree = {'a': complex(0, -1)} buff = io.BytesIO() ff = asdf.AsdfFile(tree, version="42.0.0", extensions=[FancyComplexExtension()]) ff.write_to(buff) assert b'complex-42.0.0' in buff.getvalue() del versioning._version_map['42.0.0'] versioning.supported_versions.pop()
def test_large_literals(use_numpy): largeval = 1 << 53 if use_numpy: largeval = np.uint64(largeval) tree = { 'large_int': largeval, } with pytest.raises(ValidationError): asdf.AsdfFile(tree) tree = { 'large_list': [largeval], } with pytest.raises(ValidationError): asdf.AsdfFile(tree) tree = { largeval: 'large_key', } with pytest.raises(ValidationError): asdf.AsdfFile(tree) tree = {'large_array': np.array([largeval], np.uint64)} ff = asdf.AsdfFile(tree) buff = io.BytesIO() ff.write_to(buff) ff.set_array_storage(ff.tree['large_array'], 'inline') buff = io.BytesIO() with pytest.raises(ValidationError): ff.write_to(buff) print(buff.getvalue())
def test_get_history_entries(tmpdir): """ Test edge cases for the get_history_entries API. Other cases tested above """ tmpfile = str(tmpdir.join('empty.asdf')) with asdf.AsdfFile() as af: af.write_to(tmpfile) # Make sure this works when there is no history section at all with asdf.open(tmpfile) as af: assert len(af['history']['extensions']) > 0 assert len(af.get_history_entries()) == 0
def test_input(tmpdir): tree = _get_large_tree() tmpfile = os.path.join(str(tmpdir), 'test1.asdf') afile = asdf.AsdfFile(tree) afile.write_to(tmpfile, all_array_compression='zlib') afile.close() afile = asdf.AsdfFile.open(tmpfile) tmpfile = os.path.join(str(tmpdir), 'test2.asdf') afile.write_to(tmpfile) afile.close() afile = asdf.AsdfFile.open(tmpfile) helpers.assert_tree_match(tree, afile.tree) assert afile.get_array_compression(afile.tree['science_data']) == 'zlib' afile.close()
def test_pad_blocks(tmpdir): tmpdir = str(tmpdir) # This is the case where the new tree can't fit in the available space my_array = np.ones((8, 8)) * 1 my_array2 = np.ones((42, 5)) * 2 tree = {'my_array': my_array, 'my_array2': my_array2} ff = asdf.AsdfFile(tree) ff.write_to(os.path.join(tmpdir, "test.asdf"), pad_blocks=True) with asdf.open(os.path.join(tmpdir, "test.asdf")) as ff: assert_array_equal(ff.tree['my_array'], my_array) assert_array_equal(ff.tree['my_array2'], my_array2)
def mk_flat(outfilepath): meta = {} add_ref_common(meta) flatref = stnode.FlatRef() meta['reftype'] = 'FLAT' flatref['meta'] = meta shape = (20, 20) flatref['data'] = np.zeros(shape, dtype=np.float32) flatref['dq'] = np.zeros(shape, dtype=np.uint32) flatref['err'] = np.zeros(shape, dtype=np.float32) af = asdf.AsdfFile() af.tree = {'roman': flatref} af.write_to(outfilepath)
def test_embed_asdf_in_fits_file_anonymous_extensions(tmpdir, dtype): # Write the AsdfInFits object out as a FITS file with ASDF extension asdf_in_fits = create_asdf_in_fits(dtype) asdf_in_fits.write_to(os.path.join(str(tmpdir), 'test.fits')) ff2 = asdf.AsdfFile(asdf_in_fits.tree) ff2.write_to(os.path.join(str(tmpdir), 'plain.asdf')) with fits.open(os.path.join(str(tmpdir), 'test.fits')) as hdulist: assert len(hdulist) == 4 assert [x.name for x in hdulist] == ['PRIMARY', '', '', 'ASDF'] asdf_hdu = hdulist['ASDF'] assert isinstance(asdf_hdu, fits.BinTableHDU) assert asdf_hdu.data.tobytes().startswith(b'#ASDF') with fits_embed.AsdfInFits.open(hdulist) as ff2: assert_tree_match(asdf_in_fits.tree, ff2.tree) ff = asdf.AsdfFile(copy.deepcopy(ff2.tree)) ff.write_to(os.path.join(str(tmpdir), 'test.asdf')) with asdf.open(os.path.join(str(tmpdir), 'test.asdf')) as ff: assert_tree_match(asdf_in_fits.tree, ff.tree)
def test_update_exceptions(tmpdir): tmpdir = str(tmpdir) path = os.path.join(tmpdir, 'test.asdf') my_array = np.random.rand(8, 8) tree = {'my_array': my_array} ff = asdf.AsdfFile(tree) ff.write_to(path) with asdf.open(path, mode='r', copy_arrays=True) as ff: with pytest.raises(IOError): ff.update() ff = asdf.AsdfFile(tree) buff = io.BytesIO() ff.write_to(buff) buff.seek(0) with asdf.open(buff, mode='rw') as ff: ff.update() with pytest.raises(ValueError): asdf.AsdfFile().update()
def test_extension_metadata(tmpdir): ff = asdf.AsdfFile() tmpfile = str(tmpdir.join('extension.asdf')) ff.write_to(tmpfile) with asdf.open(tmpfile) as af: assert len(af.tree['history']['extensions']) == 1 metadata = af.tree['history']['extensions'][0] assert metadata.extension_class == 'asdf.extension.BuiltinExtension' # Don't bother with testing the version here since it will depend on # how recently the package was built (version is auto-generated) assert metadata.software['name'] == 'asdf'