@@ -312,7 +312,33 @@ def _get_stub(cls: schema.Class, base_import: str, generated_import_prefix: str)
312312 else :
313313 accessors = []
314314 return ql .Stub (name = cls .name , base_import = base_import , import_prefix = generated_import_prefix ,
315- synth_accessors = accessors , ql_internal = "ql_internal" in cls .pragmas )
315+ doc = cls .doc , synth_accessors = accessors ,
316+ ql_internal = "ql_internal" in cls .pragmas )
317+
318+
319+ _stub_qldoc_header = "// the following QLdoc is generated: if you need to edit it, do it in the schema file\n "
320+
321+ _class_qldoc_re = re .compile (
322+ rf"(?P<qldoc>(?:{ re .escape (_stub_qldoc_header )} )?/\*\*.*?\*/\s*|^\s*)class\s+(?P<class>\w+)" ,
323+ re .MULTILINE | re .DOTALL )
324+
325+
326+ def _patch_class_qldoc (cls : str , qldoc : str , stub_file : pathlib .Path ):
327+ """ Replace or insert `qldoc` as the QLdoc of class `cls` in `stub_file` """
328+ if not qldoc or not stub_file .exists ():
329+ return
330+ qldoc = "\n " .join (l .rstrip () for l in qldoc .splitlines ())
331+ with open (stub_file ) as input :
332+ contents = input .read ()
333+ for match in _class_qldoc_re .finditer (contents ):
334+ if match ["class" ] == cls :
335+ qldoc_start , qldoc_end = match .span ("qldoc" )
336+ contents = f"{ contents [:qldoc_start ]} { _stub_qldoc_header } { qldoc } \n { contents [qldoc_end :]} "
337+ tmp = stub_file .with_suffix (f"{ stub_file .suffix } .bkp" )
338+ with open (tmp , "w" ) as out :
339+ out .write (contents )
340+ tmp .rename (stub_file )
341+ return
316342
317343
318344def generate (opts , renderer ):
@@ -362,9 +388,13 @@ def generate(opts, renderer):
362388 for c in data .classes .values ():
363389 path = _get_path (c )
364390 stub_file = stub_out / path
391+ base_import = get_import (out / path , opts .root_dir )
392+ stub = _get_stub (c , base_import , generated_import_prefix )
365393 if not renderer .is_customized_stub (stub_file ):
366- base_import = get_import (out / path , opts .root_dir )
367- renderer .render (_get_stub (c , base_import , generated_import_prefix ), stub_file )
394+ renderer .render (stub , stub_file )
395+ else :
396+ qldoc = renderer .render_str (stub , template = 'ql_stub_class_qldoc' )
397+ _patch_class_qldoc (c .name , qldoc , stub_file )
368398
369399 # for example path/to/elements -> path/to/elements.qll
370400 renderer .render (ql .ImportList ([i for name , i in imports .items () if not classes [name ].ql_internal ]),
0 commit comments