in python/scripts/postprocess_type_hints.py [0:0]
def main(args=None):
parser = argparse.ArgumentParser(
description="Post-process type hint files produced by pybind11-stubgen for Voyager."
)
parser.add_argument(
"source_directory",
default=os.path.join(REPO_ROOT, "pybind11-stubgen-output"),
)
parser.add_argument(
"target_directory",
default=os.path.join(REPO_ROOT, "voyager"),
)
parser.add_argument(
"--check",
action="store_true",
help=(
"Return a non-zero exit code if files on disk don't match what this script would"
" generate."
),
)
args = parser.parse_args(args)
output_file_to_source_files = defaultdict(list)
for source_path in Path(args.source_directory).rglob("*.pyi"):
output_file_to_source_files[str(source_path)].append(str(source_path))
for output_file_name, source_files in output_file_to_source_files.items():
os.makedirs(os.path.dirname(output_file_name), exist_ok=True)
print(f"Writing stub file {output_file_name}...")
file_contents = io.StringIO()
start_of_file_contents = io.StringIO()
end_of_file_contents = io.StringIO()
for source_file in source_files:
module_name = output_file_name.replace("__init__.pyi", "").replace("/", ".").rstrip(".")
with open(source_file) as f:
in_excluded_indented_block = False
in_moved_to_start_indented_block = False
in_moved_to_end_indented_block = False
for line in f:
if all(x not in line for x in OMIT_LINES_CONTAINING):
if any(line.startswith(x) for x in REMOVE_INDENTED_BLOCKS_STARTING_WITH):
in_excluded_indented_block = True
continue
elif any(line.startswith(x) for x in INDENTED_BLOCKS_TO_MOVE_TO_END):
in_moved_to_end_indented_block = True
elif any(line.startswith(x) for x in INDENTED_BLOCKS_TO_MOVE_TO_START):
in_moved_to_start_indented_block = True
elif line.strip() and not line.startswith(" "):
in_excluded_indented_block = False
in_moved_to_end_indented_block = False
if in_excluded_indented_block:
continue
for _tuple in REPLACEMENTS:
if len(_tuple) == 2:
find, replace = _tuple
only_in_module = None
else:
find, replace, only_in_module = _tuple
if only_in_module and only_in_module != module_name:
continue
results = re.findall(find, line)
if results:
line = re.sub(find, replace, line)
if in_moved_to_start_indented_block:
start_of_file_contents.write(line)
elif in_moved_to_end_indented_block:
end_of_file_contents.write(line)
else:
file_contents.write(line)
print(f"\tRead {f.tell():,} bytes of stubs from {source_file}.")
# Append end-of-file contents at the end:
file_contents.write("\n")
file_contents.write(end_of_file_contents.getvalue())
# Find the end of the __all__ block (first line with only "]"):
file_contents_string = file_contents.getvalue()
start_of_file_marker = "\n]\n"
insert_index = file_contents_string.find(start_of_file_marker)
if insert_index < 0 and len(start_of_file_contents.getvalue()):
raise NotImplementedError("Couldn't find end of __all__ block to move code blocks to!")
insert_index += len(start_of_file_marker)
file_contents_string = (
file_contents_string[:insert_index]
+ "\n"
+ start_of_file_contents.getvalue()
+ "\n"
+ file_contents_string[insert_index:]
)
# Run black:
try:
output = black.format_file_contents(
file_contents_string,
fast=False,
mode=black.FileMode(is_pyi=True, line_length=100),
)
except black.report.NothingChanged:
output = file_contents_string
if args.check:
with open(output_file_name, "r") as f:
existing = f.read()
if not stub_files_match(existing, output):
error = f"File that would be generated ({output_file_name}) "
error += "does not match existing file!\n"
error += f"Existing file had {len(existing):,} bytes, "
error += f"expected {len(output):,} bytes.\nDiff was:\n"
diff = difflib.context_diff(existing.split("\n"), output.split("\n"))
error += "\n".join([x.strip() for x in diff])
raise ValueError(error)
else:
with open(output_file_name, "w") as o:
o.write(output)
print(f"\tWrote {o.tell():,} bytes of stubs to {output_file_name}.")
print("Done!")