Migrate to ruamel.yaml and add ordered_dict parameter to `load_json…

…` and `load_yaml` plugins (#157)

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 = {}
 

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 = {}

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)

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)

requirements.txt

@@ -1,5 +1,4 @@
 colorama
-pyyaml
 jinja2
 napalm>=2.3.0
 netmiko>=2.1.1

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"

tests/plugins/tasks/data/test_data/simple.json

@@ -3,5 +3,6 @@
     "services": [
         "dhcp",
         "dns"
-    ]
+    ],
+    "a_dict": {"a": 1, "b": 2}
 }

tests/plugins/tasks/data/test_data/simple.yaml

@@ -3,3 +3,6 @@ env: test
 services:
   - dhcp
   - dns
+a_dict:
+  a: 1
+  b: 2

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)

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)