11"""Shared validation and generation helpers for the pinned Braintrust OpenAPI spec."""
22
3+ import ast
34import copy
45import difflib
56import hashlib
@@ -289,19 +290,119 @@ def _with_inline_models(
289290 return model_spec
290291
291292
292- def _single_model_module (operations : Sequence [GeneratedOperation ]) -> str :
293- tags = {operation .tag for operation in operations }
294- if len (tags ) != 1 :
295- raise CodegenError (
296- "Model generation currently requires exactly one generated OpenAPI tag; "
297- "add explicit cross-resource model partitioning before enabling another tag"
293+ _NON_MODEL_ANNOTATION_NAMES = {"Any" , "Literal" , "Mapping" , "None" , "Sequence" }
294+
295+
296+ def _operation_annotation_names (operation : GeneratedOperation ) -> Set [str ]:
297+ annotation_names : Set [str ] = set ()
298+ for type_name in [
299+ operation .request_body_type ,
300+ operation .response_type ,
301+ * (parameter .type_name for parameter in operation .parameters ),
302+ ]:
303+ if type_name :
304+ annotation_names .update (re .findall (r"\b[A-Z][A-Za-z0-9_]*\b" , type_name ))
305+ return annotation_names
306+
307+
308+ def _operation_model_roots (operations : Sequence [GeneratedOperation ]) -> Dict [str , Set [str ]]:
309+ roots : Dict [str , Set [str ]] = {}
310+ for operation in operations :
311+ roots .setdefault (operation .tag , set ()).update (
312+ _operation_annotation_names (operation ) - _NON_MODEL_ANNOTATION_NAMES
298313 )
299- return _snake_case (next (iter (tags )))
314+ return roots
315+
316+
317+ def _partition_model_source (
318+ source : str , operations : Sequence [GeneratedOperation ]
319+ ) -> Tuple [Dict [str , str ], Dict [str , str ]]:
320+ """Partition one deterministic model-generator output by resource dependency closure.
321+
322+ Definitions reached by more than one generated tag live in ``common.py``. Resource-specific
323+ modules import those shared definitions explicitly, avoiding duplicate runtime type identities.
324+ """
325+ tree = ast .parse (source )
326+ imports : List [ast .stmt ] = []
327+ definitions : List [Tuple [str , List [ast .stmt ]]] = []
328+ for node in tree .body :
329+ if isinstance (node , (ast .Import , ast .ImportFrom )):
330+ imports .append (node )
331+ continue
332+ if isinstance (node , ast .ClassDef ):
333+ names = [node .name ]
334+ elif isinstance (node , (ast .Assign , ast .AnnAssign )):
335+ targets = node .targets if isinstance (node , ast .Assign ) else [node .target ]
336+ names = [target .id for target in targets if isinstance (target , ast .Name )]
337+ elif isinstance (node , ast .Expr ) and isinstance (node .value , ast .Constant ) and isinstance (node .value .value , str ):
338+ if definitions :
339+ definitions [- 1 ][1 ].append (node )
340+ continue
341+ else :
342+ raise CodegenError (f"Unsupported generated model statement: { type (node ).__name__ } " )
343+ if len (names ) != 1 :
344+ raise CodegenError ("Generated model definitions must bind exactly one public name" )
345+ definitions .append ((names [0 ], [node ]))
346+
347+ definition_names = {name for name , _ in definitions }
348+ dependencies : Dict [str , Set [str ]] = {}
349+ for name , nodes in definitions :
350+ dependencies [name ] = {
351+ child .id
352+ for node in nodes
353+ for child in ast .walk (node )
354+ if isinstance (child , ast .Name ) and child .id in definition_names and child .id != name
355+ }
300356
357+ owners : Dict [str , Set [str ]] = {name : set () for name in definition_names }
358+ for tag , roots in _operation_model_roots (operations ).items ():
359+ pending = list (roots )
360+ seen : Set [str ] = set ()
361+ while pending :
362+ name = pending .pop ()
363+ if name in seen :
364+ continue
365+ if name not in definition_names :
366+ raise CodegenError (f"Generated resource { tag !r} references unknown model { name !r} " )
367+ seen .add (name )
368+ owners [name ].add (tag )
369+ pending .extend (dependencies [name ])
370+
371+ unreachable = sorted (name for name , tags in owners .items () if not tags )
372+ if unreachable :
373+ raise CodegenError (f"Generated models are unreachable from resource methods: { unreachable } " )
374+
375+ common_names = {name for name , tags in owners .items () if len (tags ) > 1 }
376+ module_for_name = {
377+ name : "common" if name in common_names else _snake_case (next (iter (tags ))) for name , tags in owners .items ()
378+ }
301379
302- def _model_modules (spec : Mapping [str , Any ], module : str ) -> Dict [str , str ]:
303- schemas = spec .get ("components" , {}).get ("schemas" , {})
304- return {_python_type_name (name ): module for name in schemas }
380+ def source_for (node : ast .stmt ) -> str :
381+ segment = ast .get_source_segment (source , node )
382+ if segment is None :
383+ raise CodegenError (f"Could not recover generated model source for { type (node ).__name__ } " )
384+ return segment
385+
386+ import_source = "\n " .join (source_for (node ) for node in imports )
387+ bodies : Dict [str , List [str ]] = {}
388+ for name , nodes in definitions :
389+ bodies .setdefault (module_for_name [name ], []).append ("\n " .join (source_for (node ) for node in nodes ))
390+
391+ modules : Dict [str , str ] = {}
392+ for module , blocks in sorted (bodies .items ()):
393+ common_imports = sorted (
394+ dependency
395+ for name , _ in definitions
396+ if module_for_name [name ] == module
397+ for dependency in dependencies [name ]
398+ if dependency in common_names
399+ )
400+ sections = [import_source ]
401+ if common_imports and module != "common" :
402+ sections .append (f"from .common import { ', ' .join (dict .fromkeys (common_imports ))} " )
403+ sections .append ("\n \n " .join (blocks ))
404+ modules [module ] = "\n \n " .join (section for section in sections if section ) + "\n "
405+ return modules , module_for_name
305406
306407
307408def generate_tree (output_root : Path , config : Mapping [str , Any ], spec : Mapping [str , Any ]) -> ValidationReport :
@@ -310,21 +411,29 @@ def generate_tree(output_root: Path, config: Mapping[str, Any], spec: Mapping[st
310411 operations , inline_models = _collect_generated_operations (spec , config )
311412 selected_spec = _slice_model_spec (spec , {operation .operation_id for operation in operations })
312413 model_spec = _with_inline_models (selected_spec , inline_models )
313- model_module = _single_model_module (operations )
314- model_modules = _model_modules (model_spec , model_module )
315414 output_root .mkdir (parents = True , exist_ok = True )
316415 selected_spec_path = output_root .parent / "selected-spec.json"
416+ monolithic_models_path = output_root .parent / "models.py"
317417 selected_spec_path .write_text (
318418 json .dumps (model_spec , sort_keys = True , separators = ("," , ":" ), ensure_ascii = False ) + "\n " , encoding = "utf-8"
319419 )
320420 try :
321- _generate_models (selected_spec_path , output_root / "models" / f"{ model_module } .py" , config )
421+ _generate_models (selected_spec_path , monolithic_models_path , config )
422+ model_sources , model_modules = _partition_model_source (monolithic_models_path .read_text (), operations )
322423 finally :
323424 selected_spec_path .unlink (missing_ok = True )
425+ monolithic_models_path .unlink (missing_ok = True )
426+ model_paths = []
427+ for module , body in model_sources .items ():
428+ model_path = output_root / "models" / f"{ module } .py"
429+ model_path .parent .mkdir (parents = True , exist_ok = True )
430+ _write_generated_file (model_path , body , config )
431+ model_paths .append (model_path )
324432 _write_generated_file (output_root / "__init__.py" , _GENERATED_INIT_BODY , config )
325- _write_generated_file (output_root / "models" / "__init__.py" , '"""Generated private model types."""\n ' , config )
433+ model_init_path = output_root / "models" / "__init__.py"
434+ _write_generated_file (model_init_path , _model_package_source (model_modules ), config )
326435 resource_files = _generate_resources (output_root , operations , model_modules , config )
327- _format_generated_files (resource_files )
436+ _format_generated_files ([ * model_paths , model_init_path , * resource_files ] )
328437 return report
329438
330439
@@ -459,6 +568,20 @@ def _generated_header(config: Mapping[str, Any], content_hash: str) -> str:
459568'''
460569
461570
571+ def _model_package_source (model_modules : Mapping [str , str ]) -> str :
572+ by_module : Dict [str , List [str ]] = {}
573+ for name , module in model_modules .items ():
574+ by_module .setdefault (module , []).append (name )
575+
576+ lines = ['"""Generated private model types with stable package-level imports."""' , "" ]
577+ for module , names in sorted (by_module .items ()):
578+ lines .append (f"from .{ module } import { ', ' .join (sorted (names ))} " )
579+ lines .extend (["" , "" , "__all__ = [" ])
580+ lines .extend (f" { name !r} ," for name in sorted (model_modules ))
581+ lines .extend (["]" , "" ])
582+ return "\n " .join (lines )
583+
584+
462585def _validate_selected_operations (
463586 operations : Sequence [Tuple [str , str , Any , Mapping [str , Any ], Mapping [str , Any ]]],
464587 endpoint : Mapping [str , Any ],
@@ -474,11 +597,31 @@ def _validate_selected_operations(
474597 if missing_tags :
475598 raise CodegenError (f"endpoint_generator.generated_tags contains unknown tags: { sorted (missing_tags )} " )
476599
600+ safe_reads = set (endpoint ["safe_reads" ])
601+ stale_safe_reads = safe_reads - set (supported )
602+ if stale_safe_reads :
603+ operation_id = sorted (stale_safe_reads )[0 ]
604+ raise CodegenError (f"endpoint_generator.safe_reads references non-generated operation { operation_id !r} " )
605+ non_post_safe_reads = sorted (
606+ operation_id for method , _ , operation_id , _ , _ in operations if operation_id in safe_reads and method != "post"
607+ )
608+ if non_post_safe_reads :
609+ raise CodegenError (
610+ f"endpoint_generator.safe_reads must reference POST operations; got { non_post_safe_reads [0 ]!r} "
611+ )
612+
477613 idempotent_writes = set (endpoint ["idempotent_writes" ])
478614 stale_idempotent_writes = idempotent_writes - set (supported )
479615 if stale_idempotent_writes :
480616 operation_id = sorted (stale_idempotent_writes )[0 ]
481617 raise CodegenError (f"endpoint_generator.idempotent_writes references non-generated operation { operation_id !r} " )
618+ overlapping_retry_modes = safe_reads & idempotent_writes
619+ if overlapping_retry_modes :
620+ operation_id = sorted (overlapping_retry_modes )[0 ]
621+ raise CodegenError (
622+ f"Operation { operation_id !r} cannot appear in both endpoint_generator.safe_reads "
623+ "and endpoint_generator.idempotent_writes"
624+ )
482625 non_writes = sorted (
483626 operation_id
484627 for method , _ , operation_id , _ , _ in operations
@@ -488,8 +631,8 @@ def _validate_selected_operations(
488631 raise CodegenError (f"endpoint_generator.idempotent_writes references read operation { non_writes [0 ]!r} " )
489632
490633
491- def _operation_retry_mode (method : str , operation_id : str , idempotent_writes : Set [str ]) -> str :
492- if method in {"get" , "head" }:
634+ def _operation_retry_mode (method : str , operation_id : str , safe_reads : Set [ str ], idempotent_writes : Set [str ]) -> str :
635+ if method in {"get" , "head" } or operation_id in safe_reads :
493636 return "SAFE_READ"
494637 if operation_id in idempotent_writes :
495638 return "IDEMPOTENT_WRITE"
@@ -500,6 +643,7 @@ def _collect_generated_operations(
500643 spec : Mapping [str , Any ], config : Mapping [str , Any ]
501644) -> Tuple [List [GeneratedOperation ], List [Tuple [str , Mapping [str , Any ]]]]:
502645 endpoint = _endpoint_config (config )
646+ safe_reads = set (endpoint ["safe_reads" ])
503647 idempotent_writes = set (endpoint ["idempotent_writes" ])
504648 operations : List [GeneratedOperation ] = []
505649 inline_models : Dict [str , Mapping [str , Any ]] = {}
@@ -528,7 +672,7 @@ def _collect_generated_operations(
528672 response_type = response_type ,
529673 success_statuses = statuses ,
530674 json_success_statuses = json_statuses ,
531- retry_mode = _operation_retry_mode (method , operation_id , idempotent_writes ),
675+ retry_mode = _operation_retry_mode (method , operation_id , safe_reads , idempotent_writes ),
532676 )
533677 )
534678 return operations , list (inline_models .items ())
@@ -663,18 +807,10 @@ def _generate_resources(
663807def _resource_module_source (
664808 tag : str , operations : Sequence [GeneratedOperation ], model_modules : Mapping [str , str ]
665809) -> str :
666- annotation_names : Set [str ] = set ()
667- for operation in operations :
668- for type_name in [
669- operation .request_body_type ,
670- operation .response_type ,
671- * (parameter .type_name for parameter in operation .parameters ),
672- ]:
673- if type_name :
674- annotation_names .update (re .findall (r"\b[A-Z][A-Za-z0-9_]*\b" , type_name ))
810+ annotation_names = set ().union (* (_operation_annotation_names (operation ) for operation in operations ))
675811 collections_imports = sorted (annotation_names & {"Mapping" , "Sequence" })
676812 typing_imports = sorted (annotation_names & {"Any" , "Literal" })
677- model_type_names = annotation_names - { "Any" , "Literal" , "Mapping" , "None" , "Sequence" }
813+ model_type_names = annotation_names - _NON_MODEL_ANNOTATION_NAMES
678814 model_imports : Dict [str , Set [str ]] = {}
679815 for type_name in model_type_names :
680816 module = model_modules .get (type_name )
@@ -834,13 +970,14 @@ def _endpoint_config(config: Mapping[str, Any]) -> Mapping[str, Any]:
834970 or len (generated_tags ) != len (set (generated_tags ))
835971 ):
836972 raise CodegenError ("endpoint_generator.generated_tags must be a unique list of non-empty strings" )
837- idempotent_writes = endpoint .get ("idempotent_writes" )
838- if (
839- not isinstance (idempotent_writes , list )
840- or not all (isinstance (value , str ) and value for value in idempotent_writes )
841- or len (idempotent_writes ) != len (set (idempotent_writes ))
842- ):
843- raise CodegenError ("endpoint_generator.idempotent_writes must be a unique list of non-empty strings" )
973+ for key in ("safe_reads" , "idempotent_writes" ):
974+ values = endpoint .get (key )
975+ if (
976+ not isinstance (values , list )
977+ or not all (isinstance (value , str ) and value for value in values )
978+ or len (values ) != len (set (values ))
979+ ):
980+ raise CodegenError (f"endpoint_generator.{ key } must be a unique list of non-empty strings" )
844981 for key in ("supported_request_media_types" , "supported_response_media_types" , "supported_success_statuses" ):
845982 values = endpoint .get (key )
846983 if not isinstance (values , list ) or not values or not all (isinstance (value , str ) for value in values ):
0 commit comments