2121import os
2222import shutil
2323import subprocess
24+
25+ import jinja2
2426import sys
2527import textwrap
2628from functools import lru_cache
@@ -139,18 +141,22 @@ def post_process_historically_publicised_methods(stub_file_path: Path, line: str
139141 new_lines .append (line )
140142
141143
142- def post_process_generated_stub_file (stub_file_path : Path , lines : list [str ], patch_historical_methods = False ):
144+ def post_process_generated_stub_file (
145+ module_name : str , stub_file_path : Path , lines : list [str ], patch_historical_methods = False
146+ ):
143147 """
144148 Post process the stub file - add the preamble and optionally patch historical methods.
145149 Adding preamble always, makes sure that we can update the preamble and have it automatically updated
146150 in generated files even if no API specification changes.
147151
152+ :param module_name: name of the module of the file
148153 :param stub_file_path: path of the stub fil
149154 :param lines: lines that were read from the file (with stripped comments)
150155 :param patch_historical_methods: whether we should patch historical methods
151156 :return: resulting lines of the file after post-processing
152157 """
153- new_lines = PREAMBLE .splitlines ()
158+ template = jinja2 .Template (PREAMBLE )
159+ new_lines = template .render (module_name = module_name ).splitlines ()
154160 for line in lines :
155161 if patch_historical_methods :
156162 post_process_historically_publicised_methods (stub_file_path , line , new_lines )
@@ -159,28 +165,39 @@ def post_process_generated_stub_file(stub_file_path: Path, lines: list[str], pat
159165 return new_lines
160166
161167
162- def read_pyi_file_content (pyi_file_path : Path , patch_historical_methods = False ) -> list [str ] | None :
168+ def read_pyi_file_content (
169+ module_name : str , pyi_file_path : Path , patch_historical_methods = False
170+ ) -> list [str ] | None :
163171 """
164172 Reads stub file content with post-processing and optionally patching historical methods. The comments
165- are stripped and preamble is always added. It makes sure that we can update the preamble
166- and have it automatically updated in generated files even if no API specification changes.
173+ and initial javadoc are stripped and preamble is always added. It makes sure that we can update
174+ the preamble and have it automatically updated in generated files even if no API specification changes.
167175
168176 If None is returned, the file should be deleted.
169177
170- :param pyi_file_path:
171- :param patch_historical_methods:
178+ :param module_name: name of the module in question
179+ :param pyi_file_path: the path of the file to read
180+ :param patch_historical_methods: whether the historical methods should be patched
172181 :return: list of lines of post-processed content or None if the file should be deleted.
173182 """
174- lines = [
183+ lines_no_comments = [
175184 line
176185 for line in pyi_file_path .read_text (encoding = "utf-8" ).splitlines ()
177186 if line .strip () and not line .strip ().startswith ("#" )
178187 ]
188+ remove_docstring = False
189+ lines = []
190+ for line in lines_no_comments :
191+ if line .strip ().startswith ('"""' ):
192+ remove_docstring = not remove_docstring
193+ continue
194+ if not remove_docstring :
195+ lines .append (line )
179196 if (pyi_file_path .name == "__init__.pyi" ) and lines == []:
180197 console .print (f"[yellow]Skip { pyi_file_path } )
181198 return None
182199 return post_process_generated_stub_file (
183- pyi_file_path , lines , patch_historical_methods = patch_historical_methods
200+ module_name , pyi_file_path , lines , patch_historical_methods = patch_historical_methods
184201 )
185202
186203
@@ -194,7 +211,10 @@ def compare_stub_files(generated_stub_path: Path, force_override: bool) -> tuple
194211 _removals , _additions = 0 , 0
195212 rel_path = generated_stub_path .relative_to (OUT_DIR )
196213 target_path = PROVIDERS_ROOT / rel_path
197- generated_pyi_content = read_pyi_file_content (generated_stub_path , patch_historical_methods = True )
214+ module_name = "airflow.providers." + os .fspath (rel_path .with_suffix ("" )).replace (os .path .sep , "." )
215+ generated_pyi_content = read_pyi_file_content (
216+ module_name , generated_stub_path , patch_historical_methods = True
217+ )
198218 if generated_pyi_content is None :
199219 os .unlink (generated_stub_path )
200220 if target_path .exists ():
@@ -217,7 +237,7 @@ def compare_stub_files(generated_stub_path: Path, force_override: bool) -> tuple
217237 console .print (f"[yellow]New file { target_path } )
218238 target_path .write_text ("\n " .join (generated_pyi_content ), encoding = "utf-8" )
219239 return 0 , 1
220- target_pyi_content = read_pyi_file_content (target_path , patch_historical_methods = False )
240+ target_pyi_content = read_pyi_file_content (module_name , target_path , patch_historical_methods = False )
221241 if target_pyi_content is None :
222242 target_pyi_content = []
223243 if generated_pyi_content != target_pyi_content :
@@ -288,6 +308,10 @@ def compare_stub_files(generated_stub_path: Path, force_override: bool) -> tuple
288308#
289309# You can read more in the README_API.md file
290310#
311+ \" \" \"
312+ Definition of the public interface for {{ module_name }}
313+ isort:skip_file
314+ \" \" \"
291315"""
292316
293317
0 commit comments