diff --git a/nornir/core/configuration.py b/nornir/core/configuration.py index a691343c..aa4d5f74 100644 --- a/nornir/core/configuration.py +++ b/nornir/core/configuration.py @@ -2,7 +2,7 @@ import os -import yaml +import ruamel.yaml CONF = { @@ -93,7 +93,7 @@ class Config(object): def __init__(self, config_file=None, **kwargs): if config_file: with open(config_file, "r") as f: - data = yaml.load(f.read()) or {} + data = ruamel.yaml.safe_load(f.read()) or {} else: data = {} diff --git a/nornir/plugins/inventory/simple.py b/nornir/plugins/inventory/simple.py index cf77c98a..03b87039 100644 --- a/nornir/plugins/inventory/simple.py +++ b/nornir/plugins/inventory/simple.py @@ -4,7 +4,7 @@ from nornir.core.inventory import Inventory -import yaml +import ruamel.yaml class SimpleInventory(Inventory): @@ -118,12 +118,12 @@ class SimpleInventory(Inventory): def __init__(self, host_file="hosts.yaml", group_file="groups.yaml", **kwargs): with open(host_file, "r") as f: - hosts = yaml.load(f.read()) + hosts = ruamel.yaml.safe_load(f.read()) if group_file: if os.path.exists(group_file): with open(group_file, "r") as f: - groups = yaml.load(f.read()) + groups = ruamel.yaml.safe_load(f.read()) else: logging.warning("{}: doesn't exist".format(group_file)) groups = {} diff --git a/nornir/plugins/tasks/data/load_json.py b/nornir/plugins/tasks/data/load_json.py index 06782000..453e9f6f 100644 --- a/nornir/plugins/tasks/data/load_json.py +++ b/nornir/plugins/tasks/data/load_json.py @@ -1,20 +1,33 @@ +import collections import json from nornir.core.task import Result -def load_json(task, file): +def load_json(task, file, ordered_dict=False): """ Loads a json file. Arguments: file (str): path to the file containing the json file to load + ordered_dict (bool): If set to true used OrderedDict to load maps + + Examples: + + Simple example with ``ordered_dict``:: + + > nr.run(task=load_json, + file="mydata.json", + ordered_dict=True) Returns: :obj:`nornir.core.task.Result`: * result (``dict``): dictionary with the contents of the file """ + kwargs = {} + if ordered_dict: + kwargs["object_pairs_hook"] = collections.OrderedDict with open(file, "r") as f: - data = json.loads(f.read()) + data = json.loads(f.read(), **kwargs) return Result(host=task.host, result=data) diff --git a/nornir/plugins/tasks/data/load_yaml.py b/nornir/plugins/tasks/data/load_yaml.py index 2e989097..52502419 100644 --- a/nornir/plugins/tasks/data/load_yaml.py +++ b/nornir/plugins/tasks/data/load_yaml.py @@ -1,21 +1,32 @@ from nornir.core.task import Result +import ruamel.yaml -import yaml - -def load_yaml(task, file): +def load_yaml(task, file, ordered_dict=False): """ Loads a yaml file. Arguments: file (str): path to the file containing the yaml file to load + ordered_dict (bool): If set to true used OrderedDict to load maps + + Examples: + + Simple example with ``ordered_dict``:: + + > nr.run(task=load_yaml, + file="mydata.yaml", + ordered_dict=True) Returns: :obj:`nornir.core.task.Result`: * result (``dict``): dictionary with the contents of the file """ + kwargs = {} + kwargs["typ"] = "rt" if ordered_dict else "safe" with open(file, "r") as f: - data = yaml.load(f.read()) + yml = ruamel.yaml.YAML(pure=True, **kwargs) + data = yml.load(f.read()) return Result(host=task.host, result=data) diff --git a/requirements.txt b/requirements.txt index 12396b0b..335ed1cb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ colorama -pyyaml jinja2 napalm>=2.3.0 netmiko>=2.1.1 diff --git a/tests/core/test_configuration/config.yaml b/tests/core/test_configuration/config.yaml index 164c4f87..e66cb710 100644 --- a/tests/core/test_configuration/config.yaml +++ b/tests/core/test_configuration/config.yaml @@ -1,6 +1,6 @@ --- num_workers: 10 -raise_on_error: no +raise_on_error: false user_defined: "asdasd" my_root: user_defined: "i am nested" diff --git a/tests/plugins/tasks/data/test_data/simple.json b/tests/plugins/tasks/data/test_data/simple.json index 3a17c39d..750df624 100644 --- a/tests/plugins/tasks/data/test_data/simple.json +++ b/tests/plugins/tasks/data/test_data/simple.json @@ -3,5 +3,6 @@ "services": [ "dhcp", "dns" - ] + ], + "a_dict": {"a": 1, "b": 2} } diff --git a/tests/plugins/tasks/data/test_data/simple.yaml b/tests/plugins/tasks/data/test_data/simple.yaml index fe7fdb00..3e9971f7 100644 --- a/tests/plugins/tasks/data/test_data/simple.yaml +++ b/tests/plugins/tasks/data/test_data/simple.yaml @@ -3,3 +3,6 @@ env: test services: - dhcp - dns +a_dict: + a: 1 + b: 2 diff --git a/tests/plugins/tasks/data/test_load_json.py b/tests/plugins/tasks/data/test_load_json.py index 95d8d9e3..4155478b 100644 --- a/tests/plugins/tasks/data/test_load_json.py +++ b/tests/plugins/tasks/data/test_load_json.py @@ -1,4 +1,5 @@ import os +from collections import OrderedDict from nornir.plugins.tasks import data @@ -16,6 +17,17 @@ def test_load_json(self, nornir): d = r.result assert d["env"] == "test" assert d["services"] == ["dhcp", "dns"] + assert isinstance(d["a_dict"], dict) + + def test_load_json_ordered_dict(self, nornir): + test_file = "{}/simple.json".format(data_dir) + result = nornir.run(data.load_json, file=test_file, ordered_dict=True) + + for h, r in result.items(): + d = r.result + assert d["env"] == "test" + assert d["services"] == ["dhcp", "dns"] + assert isinstance(d["a_dict"], OrderedDict) def test_load_json_error_broken_file(self, nornir): test_file = "{}/broken.json".format(data_dir) diff --git a/tests/plugins/tasks/data/test_load_yaml.py b/tests/plugins/tasks/data/test_load_yaml.py index d11b473b..9f24d6e2 100644 --- a/tests/plugins/tasks/data/test_load_yaml.py +++ b/tests/plugins/tasks/data/test_load_yaml.py @@ -1,10 +1,10 @@ import os - +from collections import OrderedDict from nornir.plugins.tasks import data -from yaml.scanner import ScannerError +from ruamel.yaml.scanner import ScannerError data_dir = "{}/test_data".format(os.path.dirname(os.path.realpath(__file__))) @@ -20,6 +20,17 @@ def test_load_yaml(self, nornir): d = r.result assert d["env"] == "test" assert d["services"] == ["dhcp", "dns"] + assert isinstance(d["a_dict"], dict) + + def test_load_yaml_ordered_dict(self, nornir): + test_file = "{}/simple.yaml".format(data_dir) + result = nornir.run(data.load_yaml, file=test_file, ordered_dict=True) + + for h, r in result.items(): + d = r.result + assert d["env"] == "test" + assert d["services"] == ["dhcp", "dns"] + assert isinstance(d["a_dict"], OrderedDict) def test_load_yaml_error_broken_file(self, nornir): test_file = "{}/broken.yaml".format(data_dir)