66from typing import (
77 Any ,
88 Dict ,
9- Optional ,
10- Protocol ,
11- Sequence ,
129 Set ,
13- Tuple ,
1410 Union ,
1511)
1612
2420FlattenedDict = Dict [str , FlattenedValue ]
2521
2622
27- class FlattenFunc (Protocol ):
28- def __call__ (
29- self ,
30- key : str ,
31- value : Any ,
32- exclude_keys : Set [str ],
33- rename_keys : Dict [str , str ],
34- flatten_functions : Dict [str , "FlattenFunc" ],
35- ** kwargs : Any ,
36- ) -> Any :
37- return None
38-
39-
4023_logger = logging .getLogger (__name__ )
4124
4225
43- def _concat_key (prefix : Optional [str ], suffix : str ):
44- if not prefix :
45- return suffix
46- return f"{ prefix } .{ suffix } "
47-
48-
49- def _is_primitive (v ):
50- for t in [str , bool , int , float ]:
51- if isinstance (v , t ):
52- return True
53- return False
54-
55-
5626def _is_homogenous_primitive_list (v ):
57- if not isinstance (v , list ):
58- return False
5927 if len (v ) == 0 :
6028 return True
61- if not _is_primitive (v [0 ]):
29+ if not isinstance (v [0 ], ( str , bool , int , float ) ):
6230 return False
6331 first_entry_value_type = type (v [0 ])
6432 for entry in v [1 :]:
@@ -67,49 +35,10 @@ def _is_homogenous_primitive_list(v):
6735 return True
6836
6937
70- def _get_flatten_func (
71- flatten_functions : Dict [str , FlattenFunc ], key_names : set [str ]
72- ) -> Optional [FlattenFunc ]:
73- for key in key_names :
74- flatten_func = flatten_functions .get (key )
75- if flatten_func is not None :
76- return flatten_func
77- return None
78-
79-
80- def _flatten_with_flatten_func (
81- key : str ,
82- value : Any ,
83- exclude_keys : Set [str ],
84- rename_keys : Dict [str , str ],
85- flatten_functions : Dict [str , FlattenFunc ],
86- key_names : Set [str ],
87- ) -> Tuple [bool , Any ]:
88- flatten_func = _get_flatten_func (flatten_functions , key_names )
89- if flatten_func is None :
90- return False , value
91- func_output = flatten_func (
92- key ,
93- value ,
94- exclude_keys = exclude_keys ,
95- rename_keys = rename_keys ,
96- flatten_functions = flatten_functions ,
97- )
98- if func_output is None :
99- return True , {}
100- if _is_primitive (func_output ) or _is_homogenous_primitive_list (
101- func_output
102- ):
103- return True , {key : func_output }
104- return False , func_output
105-
106-
10738def _flatten_compound_value_using_json (
10839 key : str ,
10940 value : Any ,
11041 exclude_keys : Set [str ],
111- rename_keys : Dict [str , str ],
112- flatten_functions : Dict [str , FlattenFunc ],
11342 _from_json = False ,
11443) -> FlattenedDict :
11544 if _from_json :
@@ -126,168 +55,64 @@ def _flatten_compound_value_using_json(
12655 value ,
12756 )
12857 return {}
129- json_value = json .loads (json_string )
13058 return _flatten_value (
13159 key ,
132- json_value ,
133- exclude_keys = exclude_keys ,
134- rename_keys = rename_keys ,
135- flatten_functions = flatten_functions ,
60+ json .loads (json_string ),
61+ exclude_keys ,
13662 # Ensure that we don't recurse indefinitely if "json.loads()" somehow returns
13763 # a complex, compound object that does not get handled by the "primitive", "list",
13864 # or "dict" cases. Prevents falling back on the JSON serialization fallback path.
139- _from_json = True ,
65+ True ,
14066 )
14167
14268
143- def _flatten_compound_value ( # pylint: disable=too-many-return-statements
69+ def _flatten_compound_value (
14470 key : str ,
14571 value : Any ,
14672 exclude_keys : Set [str ],
147- rename_keys : Dict [str , str ],
148- flatten_functions : Dict [str , FlattenFunc ],
149- key_names : Set [str ],
15073 _from_json = False ,
15174) -> FlattenedDict :
152- fully_flattened_with_flatten_func , value = _flatten_with_flatten_func (
153- key = key ,
154- value = value ,
155- exclude_keys = exclude_keys ,
156- rename_keys = rename_keys ,
157- flatten_functions = flatten_functions ,
158- key_names = key_names ,
159- )
160- if fully_flattened_with_flatten_func :
161- return value
16275 if isinstance (value , dict ):
163- return _flatten_dict (
164- value ,
165- key_prefix = key ,
166- exclude_keys = exclude_keys ,
167- rename_keys = rename_keys ,
168- flatten_functions = flatten_functions ,
169- )
76+ return flatten_dict (value , key , exclude_keys )
17077 if isinstance (value , list ):
17178 if _is_homogenous_primitive_list (value ):
17279 return {key : value }
173- return _flatten_list (
174- value ,
175- key_prefix = key ,
176- exclude_keys = exclude_keys ,
177- rename_keys = rename_keys ,
178- flatten_functions = flatten_functions ,
179- )
80+ result = {f"{ key } .length" : len (value )}
81+ for idx , val in enumerate (value ):
82+ result .update (_flatten_value (f"{ key } [{ idx } ]" , val , exclude_keys ))
83+ return result
18084 if hasattr (value , "model_dump" ):
18185 try :
182- return _flatten_dict (
183- value .model_dump (),
184- key_prefix = key ,
185- exclude_keys = exclude_keys ,
186- rename_keys = rename_keys ,
187- flatten_functions = flatten_functions ,
188- )
86+ return flatten_dict (value .model_dump (), key , exclude_keys )
18987 except TypeError :
19088 return {key : str (value )}
19189 return _flatten_compound_value_using_json (
192- key ,
193- value ,
194- exclude_keys = exclude_keys ,
195- rename_keys = rename_keys ,
196- flatten_functions = flatten_functions ,
197- _from_json = _from_json ,
90+ key , value , exclude_keys , _from_json
19891 )
19992
20093
20194def _flatten_value (
20295 key : str ,
20396 value : Any ,
20497 exclude_keys : Set [str ],
205- rename_keys : Dict [str , str ],
206- flatten_functions : Dict [str , FlattenFunc ],
20798 _from_json = False ,
20899) -> FlattenedDict :
209- if value is None :
210- return {}
211- key_names = set ([key ])
212- renamed_key = rename_keys .get (key )
213- if renamed_key is not None :
214- key_names .add (renamed_key )
215- key = renamed_key
216- if key_names & exclude_keys :
100+ if value is None or key in exclude_keys :
217101 return {}
218- if _is_primitive (value ):
102+ if isinstance (value , ( str , bool , int , float ) ):
219103 return {key : value }
220- return _flatten_compound_value (
221- key = key ,
222- value = value ,
223- exclude_keys = exclude_keys ,
224- rename_keys = rename_keys ,
225- flatten_functions = flatten_functions ,
226- key_names = key_names ,
227- _from_json = _from_json ,
228- )
104+ return _flatten_compound_value (key , value , exclude_keys , _from_json )
229105
230106
231- def _flatten_dict (
107+ def flatten_dict (
232108 d : Dict [str , Any ],
233109 key_prefix : str ,
234110 exclude_keys : Set [str ],
235- rename_keys : Dict [str , str ],
236- flatten_functions : Dict [str , FlattenFunc ],
237111) -> FlattenedDict :
238112 result = {}
239113 for key , value in d .items ():
240- if key in exclude_keys :
241- continue
242- full_key = _concat_key (key_prefix , key )
243- flattened = _flatten_value (
244- full_key ,
245- value ,
246- exclude_keys = exclude_keys ,
247- rename_keys = rename_keys ,
248- flatten_functions = flatten_functions ,
249- )
250- result .update (flattened )
251- return result
252-
253-
254- def _flatten_list (
255- lst : list [Any ],
256- key_prefix : str ,
257- exclude_keys : Set [str ],
258- rename_keys : Dict [str , str ],
259- flatten_functions : Dict [str , FlattenFunc ],
260- ) -> FlattenedDict :
261- result = {}
262- result [_concat_key (key_prefix , "length" )] = len (lst )
263- for index , value in enumerate (lst ):
264- full_key = f"{ key_prefix } [{ index } ]"
265- flattened = _flatten_value (
266- full_key ,
267- value ,
268- exclude_keys = exclude_keys ,
269- rename_keys = rename_keys ,
270- flatten_functions = flatten_functions ,
271- )
272- result .update (flattened )
114+ if key not in exclude_keys :
115+ result .update (
116+ _flatten_value (f"{ key_prefix } .{ key } " , value , exclude_keys )
117+ )
273118 return result
274-
275-
276- def flatten_dict (
277- d : Dict [str , Any ],
278- key_prefix : Optional [str ] = None ,
279- exclude_keys : Optional [Sequence [str ]] = None ,
280- rename_keys : Optional [Dict [str , str ]] = None ,
281- flatten_functions : Optional [Dict [str , FlattenFunc ]] = None ,
282- ):
283- key_prefix = key_prefix or ""
284- exclude_keys = set (exclude_keys or [])
285- rename_keys = rename_keys or {}
286- flatten_functions = flatten_functions or {}
287- return _flatten_dict (
288- d ,
289- key_prefix = key_prefix ,
290- exclude_keys = exclude_keys ,
291- rename_keys = rename_keys ,
292- flatten_functions = flatten_functions ,
293- )
0 commit comments