diff --git a/dev.py b/dev.py index f1b32767..d0edc5da 100644 --- a/dev.py +++ b/dev.py @@ -4,7 +4,7 @@ Use with `python -i dev.py` for a useful interactive shell. """ -import yaml +from ruamel import yaml from routemaster.db import * # noqa: F403, F401 from routemaster.app import App @@ -18,7 +18,7 @@ def app_from_config(config_path): By default, will use the example.yaml file. """ with open(config_path, 'r') as f: - config = load_config(yaml.load(f)) + config = load_config(yaml.safe_load(f)) class InteractiveApp(App): """ diff --git a/routemaster/cli.py b/routemaster/cli.py index 5446d87c..33f642b5 100644 --- a/routemaster/cli.py +++ b/routemaster/cli.py @@ -1,8 +1,8 @@ """CLI handling for `routemaster`.""" import logging -import yaml import click +from ruamel import yaml from routemaster.app import App from routemaster.cron import CronThread @@ -29,7 +29,7 @@ def main(ctx, config_file): logging.getLogger('schedule').setLevel(logging.CRITICAL) try: - config = load_config(yaml.load(config_file)) + config = load_config(yaml.safe_load(config_file)) except ConfigError: logger.exception("Configuration Error") click.get_current_context().exit(1) diff --git a/routemaster/config/loader.py b/routemaster/config/loader.py index db0e0883..90d01702 100644 --- a/routemaster/config/loader.py +++ b/routemaster/config/loader.py @@ -5,10 +5,10 @@ import datetime from typing import Any, Dict, List, Optional -import yaml import jsonschema import pkg_resources import jsonschema.exceptions +from ruamel import yaml from routemaster.config.model import ( Gate, @@ -71,7 +71,7 @@ def _schema_validate(config: Yaml) -> None: 'routemaster.config', 'schema.yaml', ).decode('utf-8') - schema_yaml = yaml.load(schema_raw) + schema_yaml = yaml.safe_load(schema_raw) try: jsonschema.validate(config, schema_yaml) diff --git a/routemaster/config/tests/test_loading.py b/routemaster/config/tests/test_loading.py index d35ce2ba..b7014527 100644 --- a/routemaster/config/tests/test_loading.py +++ b/routemaster/config/tests/test_loading.py @@ -5,7 +5,7 @@ from pathlib import Path from unittest import mock -import yaml +from ruamel import yaml import pytest from routemaster.config import ( @@ -37,7 +37,7 @@ def reset_environment(): def yaml_data(name: str): with open(f'test_data/{name}.yaml') as f: - return yaml.load(f) + return yaml.safe_load(f) @contextlib.contextmanager @@ -372,7 +372,7 @@ def test_example_config_loads(): assert example_yaml.exists(), "Example file is missing! (is this test set up correctly?)" - example_config = load_config(yaml.load(example_yaml.read_text())) + example_config = load_config(yaml.safe_load(example_yaml.read_text())) # Some basic assertions that we got the right thing loaded assert list(example_config.state_machines.keys()) == ['user_lifecycle'] diff --git a/routemaster/tests/test_validation.py b/routemaster/tests/test_validation.py index ff1d7601..407cef7a 100644 --- a/routemaster/tests/test_validation.py +++ b/routemaster/tests/test_validation.py @@ -1,6 +1,6 @@ from pathlib import Path -import yaml +from ruamel import yaml import pytest from routemaster.config import ( @@ -216,7 +216,7 @@ def test_example_config_is_valid(app): assert example_yaml.exists(), "Example file is missing! (is this test set up correctly?)" - example_config = load_config(yaml.load(example_yaml.read_text())) + example_config = load_config(yaml.safe_load(example_yaml.read_text())) # quick check that we've loaded the config we expect assert list(example_config.state_machines.keys()) == ['user_lifecycle'] diff --git a/setup.py b/setup.py index 5d3d10e7..22fcdfcf 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ install_requires=( 'click', - 'pyyaml', + 'ruamel.yaml', 'jsonschema >=2.6', 'flask', 'psycopg2',