diff --git a/salt/utils/dicttrim.py b/salt/utils/dicttrim.py index 58ddc83f0e80..01e8bf13d92d 100644 --- a/salt/utils/dicttrim.py +++ b/salt/utils/dicttrim.py @@ -6,6 +6,22 @@ import salt.payload +def _trim_dict_in_dict(data, max_val_size, replace_with): + ''' + Takes a dictionary, max_val_size and replace_with + and recursively loops through and replaces any values + that are greater than max_val_size. + ''' + for key in data: + if isinstance(data[key], dict): + _trim_dict_in_dict(data[key], + max_val_size, + replace_with) + else: + if sys.getsizeof(data[key]) > max_val_size: + data[key] = replace_with + + def trim_dict( data, max_dict_bytes, @@ -63,8 +79,13 @@ def trim_dict( max_val_size = float(max_dict_bytes * (percent / 100)) try: for key in data: - if sys.getsizeof(data[key]) > max_val_size: - data[key] = replace_with + if isinstance(data[key], dict): + _trim_dict_in_dict(data[key], + max_val_size, + replace_with) + else: + if sys.getsizeof(data[key]) > max_val_size: + data[key] = replace_with percent = percent - stepper_size max_val_size = float(max_dict_bytes * (percent / 100)) if use_bin_type: diff --git a/tests/unit/utils/test_dicttrim.py b/tests/unit/utils/test_dicttrim.py new file mode 100644 index 000000000000..3c21d0da4b9a --- /dev/null +++ b/tests/unit/utils/test_dicttrim.py @@ -0,0 +1,35 @@ +# -*- coding: utf-8 -*- + +# Import python libs +from __future__ import absolute_import, print_function, unicode_literals + +# Import Salt Testing libs +from tests.support.unit import TestCase + +# Import Salt libs +import salt.utils.dicttrim as dicttrimmer + +import logging +log = logging.getLogger(__name__) + + +class DictTrimTestCase(TestCase): + + def setUp(self): + self.old_dict = {'a': 'b', 'c': 'x' * 10000} + self.new_dict = {'a': 'b', 'c': 'VALUE_TRIMMED'} + + def test_trim_dict(self): + ret = dicttrimmer.trim_dict(self.old_dict, 1000) + self.assertEqual(ret, self.new_dict) + + +class RecursiveDictTrimTestCase(TestCase): + + def setUp(self): + self.old_dict = {'a': {'b': 1, 'c': 2, 'e': 'x' * 10000, 'f': '3'}} + self.new_dict = {'a': {'b': 1, 'c': 2, 'e': 'VALUE_TRIMMED', 'f': '3'}} + + def test_trim_dict(self): + ret = dicttrimmer.trim_dict(self.old_dict, 1000) + self.assertEqual(ret, self.new_dict)