add a deep_map utility function

dbt/utils.py

@@ -1,12 +1,14 @@
 from datetime import datetime
 from decimal import Decimal
-import os
-import hashlib
-import itertools
-import json
+
 import collections
 import copy
 import functools
+import hashlib
+import itertools
+import json
+import numbers
+import os
 
 import dbt.exceptions
 import dbt.flags
@@ -263,6 +265,53 @@ def deep_merge_item(destination, key, value):
         destination[key] = value
 
 
+def deep_map(func, value, keypath=(), memo=None, _notfound=object()):
+    """map the function func() onto each non-container value in 'value'
+    recursively, returning a new value. As long as func does not manipulate
+    value, then deep_map will also not manipulate it.
+
+    value should be a value returned by `yaml.safe_load` or `json.load` - the
+    only expected types are list, dict, native python number, str, NoneType,
+    and bool.
+
+    func() will be called on numbers, strings, Nones, and booleans. Its first
+    parameter will be the value, and the second will be its keypath, an
+    iterable over the __getitem__ keys needed to get to it.
+    """
+    # TODO: if we could guarantee no cycles, we would not need to memoize
+    if memo is None:
+        memo = {}
+
+    value_id = id(value)
+    cached = memo.get(value_id, _notfound)
+    if cached is not _notfound:
+        return cached
+
+    atomic_types = (int, float, basestring, type(None), bool)
+
+    if isinstance(value, list):
+        ret = [
+            deep_map(func, v, (keypath + (idx,)), memo)
+            for idx, v in enumerate(value)
+        ]
+    elif isinstance(value, dict):
+        ret = {
+            k: deep_map(func, v, (keypath + (k,)), memo)
+            for k, v in value.items()
+        }
+    elif isinstance(value, atomic_types):
+        ret = func(value, keypath)
+    else:
+        ok_types = (list, dict) + atomic_types
+        # TODO(jeb): real error
+        raise TypeError(
+            'in deep_map, expected one of {!r}, got {!r}'
+            .format(ok_types, type(value))
+        )
+    memo[value_id] = ret
+    return ret
+
+
 class AttrDict(dict):
     def __init__(self, *args, **kwargs):
         super(AttrDict, self).__init__(*args, **kwargs)

test/unit/test_utils.py

@@ -41,3 +41,90 @@ def test__simple_cases(self):
                 case['expected'], actual,
                 'failed on {} (actual {}, expected {})'.format(
                     case['description'], actual, case['expected']))
+
+
+
+class TestDeepMap(unittest.TestCase):
+    def setUp(self):
+        self.input_value = {
+            'foo': {
+                'bar': 'hello',
+                'baz': [1, 90.5, '990', '89.9'],
+            },
+            'nested': [
+                {
+                    'test': '90',
+                    'other_test': None,
+                },
+                {
+                    'test': 400,
+                    'other_test': 4.7e9,
+                },
+            ],
+        }
+
+    @staticmethod
+    def intify_all(value, _):
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            return -1
+
+    def test__simple_cases(self):
+        expected = {
+            'foo': {
+                'bar': -1,
+                'baz': [1, 90, 990, -1],
+            },
+            'nested': [
+                {
+                    'test': 90,
+                    'other_test': -1,
+                },
+                {
+                    'test': 400,
+                    'other_test': 4700000000,
+                },
+            ],
+        }
+        actual = dbt.utils.deep_map(self.intify_all, self.input_value)
+        self.assertEquals(actual, expected)
+
+        actual = dbt.utils.deep_map(self.intify_all, expected)
+        self.assertEquals(actual, expected)
+
+
+    @staticmethod
+    def special_keypath(value, keypath):
+
+        if tuple(keypath) == ('foo', 'baz', 1):
+            return 'hello'
+        else:
+            return value
+
+    def test__keypath(self):
+        expected = {
+            'foo': {
+                'bar': 'hello',
+                # the only change from input is the second entry here
+                'baz': [1, 'hello', '990', '89.9'],
+            },
+            'nested': [
+                {
+                    'test': '90',
+                    'other_test': None,
+                },
+                {
+                    'test': 400,
+                    'other_test': 4.7e9,
+                },
+            ],
+        }
+        actual = dbt.utils.deep_map(self.special_keypath, self.input_value)
+        self.assertEquals(actual, expected)
+
+        actual = dbt.utils.deep_map(self.special_keypath, expected)
+        self.assertEquals(actual, expected)
+
+
+