def apply_stub_annotations(stub_path: str, file_path: str) -> str:
with open(stub_path) as stub_file, open(file_path) as source_file:
stub = _parse(stub_file)
source = _parse(source_file)
context = CodemodContext()
if LIBCST_VERSION >= "0.3.5":
# pyre-ignore[16]: This is from the new version of libcst.
ApplyTypeAnnotationsVisitor.store_stub_in_context(context, stub)
else:
ApplyTypeAnnotationsVisitor.add_stub_to_context(context, stub)
modified_tree = ApplyTypeAnnotationsVisitor(context).transform_module(
source)
return modified_tree.code
def test_annotate_functions(self, stub: str, before: str, after: str) -> None:
context = CodemodContext()
ApplyTypeAnnotationsVisitor.add_stub_to_context(
context, parse_module(textwrap.dedent(stub.rstrip()))
)
self.assertCodemod(before, after, context_override=context)