1import sys 2 3import salt.payload 4 5 6def _trim_dict_in_dict(data, max_val_size, replace_with): 7 """ 8 Takes a dictionary, max_val_size and replace_with 9 and recursively loops through and replaces any values 10 that are greater than max_val_size. 11 """ 12 for key in data: 13 if isinstance(data[key], dict): 14 _trim_dict_in_dict(data[key], max_val_size, replace_with) 15 else: 16 if sys.getsizeof(data[key]) > max_val_size: 17 data[key] = replace_with 18 19 20def trim_dict( 21 data, 22 max_dict_bytes, 23 percent=50.0, 24 stepper_size=10, 25 replace_with="VALUE_TRIMMED", 26 is_msgpacked=False, 27 use_bin_type=False, 28): 29 """ 30 Takes a dictionary and iterates over its keys, looking for 31 large values and replacing them with a trimmed string. 32 33 If after the first pass over dictionary keys, the dictionary 34 is not sufficiently small, the stepper_size will be increased 35 and the dictionary will be rescanned. This allows for progressive 36 scanning, removing large items first and only making additional 37 passes for smaller items if necessary. 38 39 This function uses msgpack to calculate the size of the dictionary 40 in question. While this might seem like unnecessary overhead, a 41 data structure in python must be serialized in order for sys.getsizeof() 42 to accurately return the items referenced in the structure. 43 44 Ex: 45 >>> salt.utils.dicttrim.trim_dict({'a': 'b', 'c': 'x' * 10000}, 100) 46 {'a': 'b', 'c': 'VALUE_TRIMMED'} 47 48 To improve performance, it is adviseable to pass in msgpacked 49 data structures instead of raw dictionaries. If a msgpack 50 structure is passed in, it will not be unserialized unless 51 necessary. 52 53 If a msgpack is passed in, it will be repacked if necessary 54 before being returned. 55 56 :param use_bin_type: Set this to true if "is_msgpacked=True" 57 and the msgpack data has been encoded 58 with "use_bin_type=True". This also means 59 that the msgpack data should be decoded with 60 "encoding='utf-8'". 61 """ 62 if is_msgpacked: 63 dict_size = sys.getsizeof(data) 64 else: 65 dict_size = sys.getsizeof(salt.payload.dumps(data)) 66 if dict_size > max_dict_bytes: 67 if is_msgpacked: 68 if use_bin_type: 69 data = salt.payload.loads(data, encoding="utf-8") 70 else: 71 data = salt.payload.loads(data) 72 while True: 73 percent = float(percent) 74 max_val_size = float(max_dict_bytes * (percent / 100)) 75 try: 76 for key in data: 77 if isinstance(data[key], dict): 78 _trim_dict_in_dict(data[key], max_val_size, replace_with) 79 else: 80 if sys.getsizeof(data[key]) > max_val_size: 81 data[key] = replace_with 82 percent = percent - stepper_size 83 max_val_size = float(max_dict_bytes * (percent / 100)) 84 if use_bin_type: 85 dump_data = salt.payload.dumps(data, use_bin_type=True) 86 else: 87 dump_data = salt.payload.dumps(data) 88 cur_dict_size = sys.getsizeof(dump_data) 89 if cur_dict_size < max_dict_bytes: 90 if is_msgpacked: # Repack it 91 return dump_data 92 else: 93 return data 94 elif max_val_size == 0: 95 if is_msgpacked: 96 return dump_data 97 else: 98 return data 99 except ValueError: 100 pass 101 if is_msgpacked: 102 if use_bin_type: 103 return salt.payload.dumps(data, use_bin_type=True) 104 else: 105 return salt.payload.dumps(data) 106 else: 107 return data 108 else: 109 return data 110