diff options
| -rw-r--r-- | _test/test_string.py | 40 | ||||
| -rw-r--r-- | scalarstring.py | 32 |
2 files changed, 64 insertions, 8 deletions
diff --git a/_test/test_string.py b/_test/test_string.py index e6449ac..5aba3ef 100644 --- a/_test/test_string.py +++ b/_test/test_string.py @@ -173,3 +173,43 @@ class TestReplace: so = data['foo'].replace('foo', 'bar', 2) assert isinstance(so, ruamel.yaml.scalarstring.DoubleQuotedScalarString) assert so == 'bar bar bar foo' + + +class TestWalkTree: + def test_basic(self): + from ruamel.yaml.comments import CommentedMap + from ruamel.yaml.scalarstring import walk_tree + data = CommentedMap() + data[1] = 'a' + data[2] = 'with\nnewline\n' + walk_tree(data) + exp = """\ + 1: a + 2: | + with + newline + """ + assert round_trip_dump(data) == dedent(exp) + + def test_map(self): + from ruamel.yaml.compat import ordereddict + from ruamel.yaml.comments import CommentedMap + from ruamel.yaml.scalarstring import walk_tree, preserve_literal + from ruamel.yaml.scalarstring import DoubleQuotedScalarString as dq + from ruamel.yaml.scalarstring import SingleQuotedScalarString as sq + data = CommentedMap() + data[1] = 'a' + data[2] = 'with\nnew : line\n' + data[3] = '${abc}' + data[4] = 'almost:mapping' + m = ordereddict([('\n', preserve_literal), ('${', sq), (':', dq)]) + walk_tree(data, map=m) + exp = """\ + 1: a + 2: | + with + new : line + 3: '${abc}' + 4: "almost:mapping" + """ + assert round_trip_dump(data) == dedent(exp) diff --git a/scalarstring.py b/scalarstring.py index 822ae83..0bc3c47 100644 --- a/scalarstring.py +++ b/scalarstring.py @@ -62,25 +62,41 @@ def preserve_literal(s): return PreservedScalarString(s.replace('\r\n', '\n').replace('\r', '\n')) -def walk_tree(base): +def walk_tree(base, map=None): # type: (Any) -> None """ the routine here walks over a simple yaml tree (recursing in dict values and list items) and converts strings that have multiple lines to literal scalars + + You can also provide an explicit (ordered) mapping for multiple transforms + (first of which is executed): + map = ruamel.yaml.compat.ordereddict + map['\n'] = preserve_literal + map[':'] = SingleQuotedScalarString + walk_tree(data, map=map) """ - from ruamel.yaml.compat import string_types + from ruamel.yaml.compat import string_types, MutableMapping, MutableSequence + + if map is None: + map = {'\n': preserve_literal} - if isinstance(base, dict): + if isinstance(base, MutableMapping): for k in base: v = base[k] # type: Text - if isinstance(v, string_types) and '\n' in v: - base[k] = preserve_literal(v) + if isinstance(v, string_types): + for ch in map: + if ch in v: + base[k] = map[ch](v) + break else: walk_tree(v) - elif isinstance(base, list): + elif isinstance(base, MutableSequence): for idx, elem in enumerate(base): - if isinstance(elem, string_types) and '\n' in elem: # type: ignore - base[idx] = preserve_literal(elem) # type: ignore + if isinstance(elem, string_types): + for ch in map: + if ch in elem: + base[idx] = map[ch](elem) + break else: walk_tree(elem) |