diff --git a/tardis/conftest.py b/tardis/conftest.py index b3a4b858a4d..a75a07173e7 100644 --- a/tardis/conftest.py +++ b/tardis/conftest.py @@ -3,6 +3,7 @@ import tardis from tardis.atomic import AtomData from tardis.io.config_reader import Configuration +from tardis.io.util import yaml_load_config_file from astropy.tests.pytest_plugins import * # For specifying error while exception handling @@ -240,8 +241,8 @@ def included_he_atomic_data(test_data_path): @pytest.fixture def tardis_config_verysimple(): - return yaml.load( - open('tardis/io/tests/data/tardis_configv1_verysimple.yml')) + return yaml_load_config_file( + 'tardis/io/tests/data/tardis_configv1_verysimple.yml') @pytest.fixture(scope="session") diff --git a/tardis/io/config_reader.py b/tardis/io/config_reader.py index 9cc1c86ee77..fdcdeb6451b 100644 --- a/tardis/io/config_reader.py +++ b/tardis/io/config_reader.py @@ -7,13 +7,12 @@ from astropy import constants, units as u import numpy as np import pandas as pd -import yaml import tardis from tardis.io.model_reader import ( read_density_file, calculate_density_after_time, read_abundances_file) from tardis.io.config_validator import ConfigurationValidator -from tardis.io.util import YAMLLoader +from tardis.io.util import YAMLLoader, yaml_load_file from tardis import atomic from tardis.util import (species_string_to_tuple, parse_quantity, element_symbol2atomic_number, quantity_linspace) @@ -593,7 +592,7 @@ def from_yaml(cls, fname): filename or path """ try: - yaml_dict = yaml.load(file(fname)) + yaml_dict = yaml_load_file(fname) except IOError as e: logger.critical('No config file named: %s', fname) raise e @@ -625,7 +624,7 @@ def from_config_dict(cls, config_dict, config_definition_file=None): if config_definition_file is None: config_definition_file = default_config_definition_file - config_definition = yaml.load(open(config_definition_file)) + config_definition = yaml_load_file(config_definition_file) return cls(ConfigurationValidator(config_definition, config_dict).get_config()) @@ -736,9 +735,8 @@ class Configuration(ConfigurationNameSpace): @classmethod def from_yaml(cls, fname, *args, **kwargs): try: - yaml_dict = yaml.load( - open(fname), - Loader=kwargs.pop('loader', YAMLLoader)) + yaml_dict = yaml_load_file(fname, + loader=kwargs.pop('loader', YAMLLoader)) except IOError as e: logger.critical('No config file named: %s', fname) raise e @@ -793,7 +791,7 @@ def from_config_dict(cls, config_dict, atom_data=None, test_parser=False, if config_definition_file is None: config_definition_file = default_config_definition_file - config_definition = yaml.load(open(config_definition_file)) + config_definition = yaml_load_file(config_definition_file) if validate: validated_config_dict = ConfigurationValidator(config_definition, config_dict).get_config() diff --git a/tardis/io/tests/test_config_reader.py b/tardis/io/tests/test_config_reader.py index 7a51ddb3d03..6f04b3f19ff 100644 --- a/tardis/io/tests/test_config_reader.py +++ b/tardis/io/tests/test_config_reader.py @@ -7,8 +7,8 @@ import numpy as np from numpy.testing import assert_almost_equal, assert_array_almost_equal -from tardis.util import parse_quantity -from tardis.io.util import YAMLLoader +from tardis.io.util import YAMLLoader, yaml_load_config_file, quantity_from_str + def data_path(filename): data_dir = os.path.dirname(__file__) @@ -80,18 +80,17 @@ def setup(self): #general parsing of the paper config self.config = config_reader.Configuration.from_yaml(data_path('paper1_tardis_configv1.yml'), test_parser=True) - self.yaml_data = yaml.load(open(data_path('paper1_tardis_configv1.yml'))) - - + self.yaml_data = yaml_load_config_file( + data_path('paper1_tardis_configv1.yml')) def test_abundances(self): oxygen_abundance = self.yaml_data['model']['abundances']['O'] assert_array_almost_equal(oxygen_abundance, self.config.abundances.ix[8].values) def test_velocities(self): - assert_almost_equal(parse_quantity(self.yaml_data['model']['structure']['velocity']['start']).cgs.value, + assert_almost_equal(self.yaml_data['model']['structure']['velocity']['start'].cgs.value, self.config.structure.v_inner[0].cgs.value) - assert_almost_equal(parse_quantity(self.yaml_data['model']['structure']['velocity']['stop']).cgs.value, + assert_almost_equal(self.yaml_data['model']['structure']['velocity']['stop'].cgs.value, self.config.structure.v_outer[-1].cgs.value) assert len(self.config.structure.v_outer) == (self.yaml_data['model']['structure']['velocity']['num']) @@ -118,9 +117,9 @@ def test_number_of_packets(self): def test_spectrum_section(self): assert_almost_equal(self.config['spectrum']['start'].value, - parse_quantity(self.yaml_data['spectrum']['start']).value) + self.yaml_data['spectrum']['start'].value) assert_almost_equal(self.config['spectrum']['end'].value, - parse_quantity(self.yaml_data['spectrum']['stop']).value) + self.yaml_data['spectrum']['stop'].value) assert self.config['spectrum']['bins'] == self.yaml_data['spectrum']['num'] @@ -132,7 +131,7 @@ def test_time_explosion(self): def test_last_no_of_packets(): - yaml_data = yaml.load(open(data_path('paper1_tardis_configv1.yml'))) + yaml_data = yaml_load_config_file(data_path('paper1_tardis_configv1.yml')) del yaml_data['montecarlo']['last_no_of_packets'] config = config_reader.Configuration.from_config_dict(yaml_data, test_parser=True, @@ -147,7 +146,7 @@ def setup(self): filename = 'tardis_configv1_ascii_density.yml' self.config = config_reader.Configuration.from_yaml(data_path(filename), test_parser=True) - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) def test_velocities(self): @@ -166,7 +165,7 @@ def setup(self): filename = 'tardis_configv1_artis_density.yml' self.config = config_reader.Configuration.from_yaml(data_path(filename), test_parser=True) - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) def test_velocities(self): @@ -186,7 +185,7 @@ def setup(self): #general parsing of the paper config filename = 'tardis_configv1_artis_density.yml' - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) self.yaml_data['model']['abundances'] = {'type': 'file', 'filename': 'artis_abundances.dat', 'filetype': 'artis'} @@ -210,7 +209,7 @@ def setup(self): #general parsing of the paper config filename = 'tardis_configv1_artis_density_v_slice.yml' - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) self.yaml_data['model']['abundances'] = {'type': 'file', 'filename': 'artis_abundances.dat', 'filetype': 'artis'} @@ -235,7 +234,7 @@ def setup(self): #general parsing of the paper config filename = 'tardis_configv1_uniform_density.yml' - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) self.config = config_reader.Configuration.from_config_dict(self.yaml_data, test_parser=True, @@ -251,8 +250,8 @@ def setup(self): #general parsing of the paper config filename = 'tardis_configv1_uniform_density.yml' - self.yaml_data = yaml.load(open(data_path(filename))) - self.yaml_data['plasma']['initial_t_inner'] = "2508 K" + self.yaml_data = yaml_load_config_file(data_path(filename)) + self.yaml_data['plasma']['initial_t_inner'] = quantity_from_str("2508 K") self.config = config_reader.Configuration.from_config_dict(self.yaml_data, test_parser=True, @@ -268,7 +267,7 @@ def setup(self): #general parsing of the paper config filename = 'tardis_configv1_ascii_density_abund.yml' - self.yaml_data = yaml.load(open(data_path(filename))) + self.yaml_data = yaml_load_config_file(data_path(filename)) self.yaml_data['model']['structure']['filename'] = 'density.dat' self.yaml_data['model']['abundances']['filename'] = 'abund.dat' @@ -309,8 +308,8 @@ def test_densities(self): def test_ascii_reader_power_law(): - with open(data_path('tardis_configv1_density_power_law_test.yml')) as f: - yaml_data = yaml.load(f) + yaml_data = yaml_load_config_file( + data_path('tardis_configv1_density_power_law_test.yml')) #for later use density_data = yaml_data['model']['structure']['density'] t_explosion = density_data['time_0'] @@ -337,8 +336,8 @@ def test_ascii_reader_power_law(): def test_ascii_reader_exponential_law(): - with open(data_path('tardis_configv1_density_exponential_test.yml')) as f: - yaml_data = yaml.load(f) + yaml_data = yaml_load_config_file( + data_path('tardis_configv1_density_exponential_test.yml')) #for later use density_data = yaml_data['model']['structure']['density'] t_explosion = density_data['time_0'] diff --git a/tardis/io/util.py b/tardis/io/util.py index a51879857fc..3d7a4007ac5 100644 --- a/tardis/io/util.py +++ b/tardis/io/util.py @@ -129,6 +129,12 @@ def mapping_constructor(self, node): YAMLLoader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, YAMLLoader.mapping_constructor) +def yaml_load_file(filename, loader=yaml.Loader): + with open(filename) as stream: + return yaml.load(stream, loader) + +def yaml_load_config_file(filename): + return yaml_load_file(filename, YAMLLoader) def parse_abundance_dict_to_dataframe(abundance_dict): atomic_number_dict = dict([(element_symbol2atomic_number(symbol), abundance_dict[symbol]) diff --git a/tardis/plasma/tests/test_plasmas_full.py b/tardis/plasma/tests/test_plasmas_full.py index 11fa2ab596d..9e9f8dbc633 100644 --- a/tardis/plasma/tests/test_plasmas_full.py +++ b/tardis/plasma/tests/test_plasmas_full.py @@ -1,15 +1,14 @@ import pytest import numpy as np -import yaml import tardis import numpy.testing as nptesting from astropy import units as u import os import h5py - - from tardis.base import run_tardis +from tardis.io.util import yaml_load_config_file + def data_path(fname): return os.path.join(tardis.__path__[0], 'plasma', 'tests', 'data', fname) @@ -34,12 +33,12 @@ def setup(self): " does not seem to " "exist".format( self.atom_data_filename)) - self.config_yaml = yaml.load(open( - 'tardis/plasma/tests/data/plasma_test_config_lte.yml')) + self.config_yaml = yaml_load_config_file( + 'tardis/plasma/tests/data/plasma_test_config_lte.yml') self.config_yaml['atom_data'] = self.atom_data_filename self.lte_model = run_tardis(self.config_yaml) - self.config_yaml = yaml.load(open( - 'tardis/plasma/tests/data/plasma_test_config_nlte.yml')) + self.config_yaml = yaml_load_config_file( + 'tardis/plasma/tests/data/plasma_test_config_nlte.yml') self.config_yaml['atom_data'] = self.atom_data_filename self.nlte_model = run_tardis(self.config_yaml) diff --git a/tardis/tests/test_tardis_full.py b/tardis/tests/test_tardis_full.py index bb8f9345cf3..2ac0b00c5e2 100644 --- a/tardis/tests/test_tardis_full.py +++ b/tardis/tests/test_tardis_full.py @@ -1,11 +1,11 @@ import pytest import numpy as np -import yaml import tardis import numpy.testing as nptesting from astropy import units as u import os +from tardis.io.util import yaml_load_config_file from tardis.simulation.base import Simulation from tardis.model import Radial1DModel @@ -34,8 +34,8 @@ def setup(self): " does not seem to " "exist".format( self.atom_data_filename)) - self.config_yaml = yaml.load(open( - 'tardis/io/tests/data/tardis_configv1_verysimple.yml')) + self.config_yaml = yaml_load_config_file( + 'tardis/io/tests/data/tardis_configv1_verysimple.yml') self.config_yaml['atom_data'] = self.atom_data_filename tardis_config = Configuration.from_config_dict(self.config_yaml) diff --git a/tardis/tests/tests_slow/test_w7.py b/tardis/tests/tests_slow/test_w7.py index d8b720eae16..48cbb8d047d 100644 --- a/tardis/tests/tests_slow/test_w7.py +++ b/tardis/tests/tests_slow/test_w7.py @@ -1,11 +1,11 @@ import os -import yaml import pytest import matplotlib.pyplot as plt from numpy.testing import assert_allclose from astropy.tests.helper import assert_quantity_allclose from tardis.atomic import AtomData +from tardis.io.util import yaml_load_config_file from tardis.simulation.base import Simulation from tardis.model import Radial1DModel from tardis.io.config_reader import Configuration @@ -33,7 +33,7 @@ def setup(self, request, reference, data_path, atomic_data_fname, # data seperately and provide it to tardis_config later. For rest of # the two, we form dictionary from the config file and override those # parameters by putting file paths of these two files at proper places. - config_yaml = yaml.load(open(self.config_file)) + config_yaml = yaml_load_config_file(self.config_file) config_yaml['model']['abundances']['filename'] = self.abundances config_yaml['model']['structure']['filename'] = self.densities diff --git a/tardis/util.py b/tardis/util.py index a37e306a414..d442007b4de 100644 --- a/tardis/util.py +++ b/tardis/util.py @@ -199,7 +199,8 @@ def create_synpp_yaml(radial1d_mdl, fname, shell_no=0, lines_db=None): relevant_synpp_refs = radial1d_mdl.atom_data.synpp_refs[radial1d_mdl.atom_data.synpp_refs['ref_log_tau'] > -50] - yaml_reference = yaml.load(file(synpp_default_yaml_fname)) + with open(synpp_default_yaml_fname) as stream: + yaml_reference = yaml.load(stream) if lines_db is not None: yaml_reference['opacity']['line_dir'] = os.path.join(lines_db, 'lines') @@ -234,8 +235,8 @@ def create_synpp_yaml(radial1d_mdl, fname, shell_no=0, lines_db=None): yaml_setup['v_min'].append(yaml_reference['opacity']['v_ref']) yaml_setup['v_max'].append(yaml_reference['grid']['v_outer_max']) yaml_setup['aux'].append(1e200) - - yaml.dump(yaml_reference, stream=file(fname, 'w'), explicit_start=True) + with open(fname, 'w') as f: + yaml.dump(yaml_reference, stream=f, explicit_start=True) def intensity_black_body(nu, T):