Merge pull request #575 from ftsamis/use-yamlloader

Replace misused calls to yaml.load with a helper function.

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")

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()

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']

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])

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)
 

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)

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
 

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):