def run_test_case_with_flags(
self,
stub: str,
before: str,
after: str,
**kwargs: bool,
) -> None:
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context, parse_module(textwrap.dedent(stub.rstrip())))
# Test setting the flag on the codemod instance.
# pyre-fixme[6]: Expected `Optional[typing.Sequence[str]]` for 4th param but
# got `Dict[str, bool]`.
# pyre-fixme[6]: Expected `Optional[str]` for 4th param but got `Dict[str,
# bool]`.
# pyre-fixme[6]: Expected `bool` for 4th param but got `Dict[str, bool]`.
self.assertCodemod(before, after, context_override=context, **kwargs)
# Test setting the flag when storing the stub in the context.
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context,
parse_module(textwrap.dedent(stub.rstrip())),
**kwargs,
)
self.assertCodemod(before, after, context_override=context)
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()
ApplyTypeAnnotationsVisitor.store_stub_in_context(context, stub)
modified_tree = ApplyTypeAnnotationsVisitor(context).transform_module(source)
return modified_tree.code
def run_simple_test_case(
self,
stub: str,
before: str,
after: str,
) -> None:
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context, parse_module(textwrap.dedent(stub.rstrip())))
self.assertCodemod(before, after, context_override=context)
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_with_existing_annotations(
self, stub: str, before: str, after: str) -> None:
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context, parse_module(textwrap.dedent(stub.rstrip())))
# Test setting the overwrite flag on the codemod instance.
self.assertCodemod(before,
after,
context_override=context,
overwrite_existing_annotations=True)
# Test setting the flag when storing the stub in the context.
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context,
parse_module(textwrap.dedent(stub.rstrip())),
overwrite_existing_annotations=True,
)
self.assertCodemod(before, after, context_override=context)
def run_test_case_with_flags(
self,
stub: str,
before: str,
after: str,
**kwargs: Dict[str, bool],
) -> None:
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context, parse_module(textwrap.dedent(stub.rstrip())))
# Test setting the flag on the codemod instance.
self.assertCodemod(before, after, context_override=context, **kwargs)
# Test setting the flag when storing the stub in the context.
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context,
parse_module(textwrap.dedent(stub.rstrip())),
**kwargs,
)
self.assertCodemod(before, after, context_override=context)
def test_count_annotations(
self,
stub: str,
before: str,
after: str,
annotation_counts: AnnotationCounts,
any_changes_applied: False,
):
stub = self.make_fixture_data(stub)
before = self.make_fixture_data(before)
after = self.make_fixture_data(after)
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context=context, stub=parse_module(stub))
visitor = ApplyTypeAnnotationsVisitor(context=context)
output_code = visitor.transform_module(parse_module(before)).code
self.assertEqual(after, output_code)
self.assertEqual(str(annotation_counts),
str(visitor.annotation_counts))
self.assertEqual(any_changes_applied,
visitor.annotation_counts.any_changes_applied())
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)
