added multilevel get on CommentedMap (mlget())

5 files changed, 89 insertions, 3 deletions

diff --git a/README.rst b/README.rst
index bd6601b..0ea4283 100644
--- a/README.rst
+++ b/README.rst
@@ -14,7 +14,8 @@ Major differences with PyYAML 3.11:
 - support for simple lists as mapping keys by transformation to tuples
 - ``!!omap`` generates ordereddict (C) on Python 2, collections.OrderedDict
   on Python 3, and ``!!omap`` is generated for these types.
-- some YAML 1.2 enhancements (``0o`` octal prefix, ``\/`` escape)
+- some `YAML 1.2 <http://yaml.org/spec/1.2/spec.html>`_ enhancements 
+  (``0o`` octal prefix, ``\/`` escape)
 - pep8 compliance
 - tox and py.test based testing
 - Tests whether the C yaml library is installed as well as the header
@@ -140,6 +141,29 @@ eventuallly resulting in a different Python object (subclass or alternative),
 that should behave like the original, but on the way from Python to YAML
 generates the original (or at least something much closer).
 
+Smartening
+==========
+
+When you use round-tripping, then the complex data you get are 
+already subclasses of the build in types. So you can patch
+in extra methods or override existing ones. Some methods are already
+included and you can do::
+
+    yaml_str = """\
+    a:
+    - b:
+      c: 42
+    - d:
+        f: 196
+      e:
+        g: 3.14
+    """
+    
+    
+    data = yaml.load(yaml_str, Loader=yaml.RoundTripLoader)
+    
+    assert data.mlget(['a', 1, 'd', 'f'], list_ok=True) == 196
+    
 
 Examples
 ========
diff --git a/py/__init__.py b/py/__init__.py
index 6449209..3e4e760 100644
--- a/py/__init__.py
+++ b/py/__init__.py
@@ -21,7 +21,7 @@ def _convert_version(tup):
     return ret_val
 
 
-version_info = (0, 9, 6)
+version_info = (0, 9, 7)
 __version__ = _convert_version(version_info)
 
 del _convert_version
diff --git a/py/comments.py b/py/comments.py
index 0d83d14..6003f5b 100644
--- a/py/comments.py
+++ b/py/comments.py
@@ -259,6 +259,29 @@ class CommentedMap(ordereddict, CommentedBase):
             for x in vals:
                 self[x] = vals[x]
 
+    def mlget(self, key, default=None, list_ok=False):
+        """multi-level get that expects dicts within dicts"""
+        if not isinstance(key, list):
+            return self.get(key, default)
+        # assume that the key is a list of recursively accessible dicts
+        def get_one_level(key_list, level, d):
+            if not list_ok:
+                assert isinstance(d, dict)
+            if level >= len(key_list):
+                if level > len(key_list):
+                    raise IndexError
+                return d[key_list[level-1]]
+            return get_one_level(key_list, level+1, d[key_list[level-1]])
+
+        try:
+            return get_one_level(key, 1, self)
+        except KeyError:
+            return default
+        except (TypeError, IndexError):
+            if not list_ok:
+                Raise
+            return default
+
 
 class CommentedOrderedMap(CommentedMap):
     __slots__ = [Comment.attrib, ]
diff --git a/py/yaml.py b/py/yaml.py
index 3a2e669..55ba187 100644
--- a/py/yaml.py
+++ b/py/yaml.py
@@ -134,11 +134,32 @@ class YAML:
             print('---', repr(self.first_non_empty_line(inp)))
             print('<<<', repr(self.last_non_empty_line(inp)))
 
-        if True:
+        if False:
             for x in self._args.file:
                 scan_file(x)
             return
 
+        if True:
+            import pickle
+            lines = 0
+            for x in self._args.file:
+                print(x, end=' ')
+                if x.endswith('.yaml'):
+                    data = ruamel.yaml.load(open(x))
+                    print(len(data), end=' ')
+                    lines += len(data)
+                    out_name = x.replace('.yaml', '.pickle')
+                    with open(out_name, 'w') as fp:
+                        pickle.dump(data, fp)
+                elif x.endswith('.pickle'):
+                    with open(x) as fp:
+                        data = pickle.load(fp)
+                    print(len(data), end=' ')
+                    lines += len(data)
+                print()
+            print('lines', lines)
+            return
+
         input = dedent("""
         application: web2py
         version: 1
diff --git a/test/test_comments.py b/test/test_comments.py
index 6a400b5..f9e140f 100644
--- a/test/test_comments.py
+++ b/test/test_comments.py
@@ -226,3 +226,21 @@ class TestComments:
         # last one
         - d: 4
         """)
+
+
+class TestMultiLevelGet:
+    def test_mlget_00(self):
+        x = dedent("""\
+        a:
+        - b:
+          c: 42
+        - d:
+            f: 196
+          e:
+            g: 3.14
+        """)
+        d = round_trip_load(x)
+        assert d.mlget(['a', 1, 'd', 'f'], list_ok=True) == 196
+        with pytest.raises(AssertionError):
+            d.mlget(['a', 1, 'd', 'f']) == 196
+