def test_zip_nested(self): """Test the zip nested function.""" arg0 = { 'a': 1, 'b': { 'c': 2, 'e': [3, 4, 5], }, } arg1 = { 'a': 10, 'b': { 'c': 20, 'e': [30, 40, 50], }, } result = py_utils.zip_nested(arg0, arg1) self.assertEqual( result, { 'a': (1, 10), 'b': { 'c': (2, 20), 'e': [(3, 30), (4, 40), (5, 50)], }, }) result = py_utils.zip_nested(1, 2) self.assertEqual(result, (1, 2))
def test_dict_only(self): def map_fn(x): return x[0] + x[1] arg0 = { 'a': (1, 2), 'b': { 'c': 2, 'e': [3, 4, 5], }, } arg1 = { 'a': (10, 20), 'b': { 'c': 20, 'e': [30, 40, 50], }, } result = py_utils.zip_nested(arg0, arg1, dict_only=True) self.assertEqual( result, { 'a': ((1, 2), (10, 20)), 'b': { 'c': (2, 20), 'e': ([3, 4, 5], [30, 40, 50]), }, }) result = py_utils.map_nested(map_fn, result, dict_only=True) self.assertEqual(result, { 'a': (1, 2, 10, 20), 'b': { 'c': 22, 'e': [3, 4, 5, 30, 40, 50], }, })