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 test_annotate_functions_with_existing_annotations(
self, stub: str, before: str, after: str) -> None:
context = CodemodContext()
ApplyStubberAnnotationsVisitor.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,
python_version=PYTHON_GRAMMER,
)
# Test setting the flag when storing the stub in the context.
context = CodemodContext()
ApplyStubberAnnotationsVisitor.store_stub_in_context(
context,
parse_module(textwrap.dedent(stub.rstrip())),
overwrite_existing_annotations=True,
)
self.assertCodemod(
before,
after,
context_override=context,
python_version=PYTHON_GRAMMER,
)
def test_insert(self) -> None:
"""
Should add assignments before and after the x = ...
"""
before = """
x = 1
if True:
x = 1
"""
expected_after = """
y = 1
x = 1
z = 1
if True:
y = 1
x = 1
z = 1
"""
actual_after = self.insert_statements(
InsertAssignAroundIntegerVisitor(CodemodContext()), before
)
self.assertCodeEqual(expected_after, actual_after)
def test_leaky_codemod(self) -> None:
with temp_workspace() as tmp:
# File to trigger codemod
example: Path = tmp / "example.py"
example.write_text("""print("Hello")""")
# File that should not be modified
other = tmp / "other.py"
other.touch()
# Run command
command_instance = PrintToPPrintCommand(CodemodContext())
files = gather_files(".")
result = parallel_exec_transform_with_prettyprint(
command_instance,
files,
format_code=False,
hide_progress=True,
)
# Check results
self.assertEqual(2, result.successes)
self.assertEqual(0, result.skips)
self.assertEqual(0, result.failures)
# Expect example.py to be modified
self.assertIn(
"from pprint import pprint",
example.read_text(),
"import missing in example.py",
)
# Expect other.py to NOT be modified
self.assertNotIn(
"from pprint import pprint",
other.read_text(),
"import found in other.py",
)
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
"""
Collect type annotations from all stubs and apply them to ``tree``.
Gather existing imports from ``tree`` so that we don't add duplicate imports.
"""
import_gatherer = GatherImportsVisitor(CodemodContext())
tree.visit(import_gatherer)
existing_import_names = _get_import_names(import_gatherer.all_imports)
context_contents = self.context.scratch.get(
ApplyTypeAnnotationsVisitor.CONTEXT_KEY)
if context_contents is not None:
stub, overwrite_existing_annotations = context_contents
self.overwrite_existing_annotations = (
self.overwrite_existing_annotations
or overwrite_existing_annotations)
visitor = TypeCollector(existing_import_names, self.context)
stub.visit(visitor)
self.annotations.function_annotations.update(
visitor.function_annotations)
self.annotations.attribute_annotations.update(
visitor.attribute_annotations)
self.annotations.class_definitions.update(
visitor.class_definitions)
tree_with_imports = AddImportsVisitor(
self.context).transform_module(tree)
return tree_with_imports.visit(self)
def test_dont_add_relative_object_simple(self) -> None:
"""
Should not add object as an import since it exists.
"""
before = """
from .c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from .c import D
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[("a.b.c", "D", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def run(
removed_in: List[Tuple[int, int]],
deprecated_in: List[Tuple[int, int]],
codemod: List[str],
src: Tuple[str, ...],
) -> None:
"""
Automatically fixes deprecations removed Django deprecations.
This command takes the path to target as argument and a version of
Django to select code modifications to apply.
"""
codemodders_set = set()
for version in removed_in:
codemodders_set |= set(REMOVED_IN[version])
for version in deprecated_in:
codemodders_set |= set(DEPRECATED_IN[version])
for name in codemod:
codemodders_set.add(BY_NAME[name])
if not codemodders_set:
raise click.UsageError(
"No codemods were selected. "
"Specify '--removed-in' and/or '--deprecated-in' and/or '--codemod'."
)
codemodders_list = sorted(codemodders_set, key=lambda m: m.__name__)
click.echo(
f"Running codemods: {', '.join(m.__name__ for m in codemodders_list)}")
command_instance = BaseCodemodCommand(codemodders_list, CodemodContext())
files = get_sources(src)
call_command(command_instance, files)
def test_add_object_relative_modify_simple(self) -> None:
"""
Should modify existing import to add new object
"""
before = """
from .c import E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from .c import D, E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[("a.b.c", "D", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def build_command(codemodders_list: List) -> BaseCodemodCommand:
"""Build a custom command with the list of visitors."""
class CustomCommand(BaseCodemodCommand):
transformers = codemodders_list
return CustomCommand(CodemodContext())
def test_dont_remove_wrong_importfrom_relative(self) -> None:
"""
Should not remove import from which is relative since it is the wrong module.
"""
before = """
import bar
from .c import qux
def foo() -> None:
pass
"""
after = """
import bar
from .c import qux
def foo() -> None:
pass
"""
self.assertCodemod(
before,
after,
[("a.b.d", "qux", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_remove_importfrom_relative(self) -> None:
"""
Should remove import from which is relative
"""
before = """
import bar
from .c import qux
def foo() -> None:
pass
"""
after = """
import bar
def foo() -> None:
pass
"""
self.assertCodemod(
before,
after,
[("a.b.c", "qux", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def test_runner_default(self) -> None:
before = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
# A comment
def foo() -> None:
pass
def bar() -> int:
return 5
"""
class SimpleCodemod(Codemod):
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
self.warn("Testing")
return tree.with_changes(
header=[cst.EmptyLine(comment=cst.Comment("# A comment"))])
transform = SimpleCodemod(CodemodContext())
response = transform_module(transform, dedent(before))
self.assertIsInstance(response, TransformSuccess)
assert isinstance(response, TransformSuccess)
self.assertCodeEqual(response.code, after)
self.assertEqual(response.warning_messages, ["Testing"])
def test_remove_import_with_all(self) -> None:
"""
Make sure that if an import node itself is requested for
removal, we don't remove it if it shows up in an __all__
node.
"""
before = """
from foo import bar
from qux import baz
__all__ = ["baz"]
"""
after = """
from qux import baz
__all__ = ["baz"]
"""
class RemoveImportTransformer(VisitorBasedCodemodCommand):
METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
RemoveImportsVisitor.remove_unused_import_by_node(
self.context, node)
module = cst.parse_module(self.make_fixture_data(before))
self.assertCodeEqual(
after,
RemoveImportTransformer(
CodemodContext()).transform_module(module).code,
)
def test_add_object_resolve_dotted_relative_modify_simple(self) -> None:
"""
Should merge a relative new module with an absolute existing one.
"""
before = """
from ..c import E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
after = """
from ..c import D, E, F
def foo() -> None:
pass
def bar() -> int:
return 5
"""
self.assertCodemod(
before,
after,
[ImportItem("..c", "D", None)],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def gather_imports(self, code: str) -> Set[str]:
mod = MetadataWrapper(parse_module(
CodemodTest.make_fixture_data(code)))
mod.resolve_many(GatherUnusedImportsVisitor.METADATA_DEPENDENCIES)
instance = GatherUnusedImportsVisitor(CodemodContext(wrapper=mod))
mod.visit(instance)
return set(alias.evaluated_alias or alias.evaluated_name
for alias, _ in instance.unused_imports)
def gather_comments(self, code: str) -> GatherCommentsVisitor:
mod = MetadataWrapper(parse_module(
CodemodTest.make_fixture_data(code)))
mod.resolve_many(GatherCommentsVisitor.METADATA_DEPENDENCIES)
instance = GatherCommentsVisitor(CodemodContext(wrapper=mod),
r".*\Wnoqa(\W.*)?$")
mod.visit(instance)
return instance
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 gather_names(self,
code: str) -> GatherNamesFromStringAnnotationsVisitor:
mod = MetadataWrapper(parse_module(
CodemodTest.make_fixture_data(code)))
mod.resolve_many(
GatherNamesFromStringAnnotationsVisitor.METADATA_DEPENDENCIES)
instance = GatherNamesFromStringAnnotationsVisitor(
CodemodContext(wrapper=mod))
mod.visit(instance)
return instance
def apply_stub_using_libcst(stub: str, source: str) -> str:
try:
stub_module = parse_module(stub)
source_module = parse_module(source)
context = CodemodContext()
ApplyTypeAnnotationsVisitor.add_stub_to_context(context, stub_module)
transformer = ApplyTypeAnnotationsVisitor(context)
transformed_source_module = transformer.transform_module(source_module)
except Exception as exception:
raise HandlerError(f"Failed applying stub with libcst:\n{exception}")
return transformed_source_module.code
def _import_annotations_from_future(self) -> None:
"""We need this because the original sqlalchemy types aren't generic
and will fail at runtime."""
LOG.info("Importing necessary annotations...")
context = CodemodContext()
AddImportsVisitor.add_needed_import(context, "__future__",
"annotations")
for path in self.paths:
source = libcst.parse_module(path.read_text())
modified_tree = AddImportsVisitor(context).transform_module(source)
path.write_text(modified_tree.code)
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 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()
ApplyStubberAnnotationsVisitor.store_stub_in_context(
context,
stub=parse_module(textwrap.dedent(stub.rstrip())),
overwrite_existing_annotations=True,
)
self.assertCodemod(
before,
after,
context_override=context,
python_version=PYTHON_GRAMMER,
)
def test_runner_interrupted(self) -> None:
code = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
class SimpleCodemod(Codemod):
def transform_module_impl(self, tree: cst.Module) -> cst.Module:
raise KeyboardInterrupt("Testing")
transform = SimpleCodemod(CodemodContext())
response = transform_module(transform, dedent(code))
self.assertIsInstance(response, TransformExit)
def test_remove_import_alias_after_inserting(self) -> None:
before = "from foo import bar, baz"
after = "from foo import quux, baz"
class AddRemoveTransformer(VisitorBasedCodemodCommand):
def visit_Module(self, node: cst.Module) -> None:
AddImportsVisitor.add_needed_import(self.context, "foo",
"quux")
RemoveImportsVisitor.remove_unused_import(
self.context, "foo", "bar")
module = cst.parse_module(self.make_fixture_data(before))
self.assertCodeEqual(
AddRemoveTransformer(
CodemodContext()).transform_module(module).code,
after,
)
def test_noop(self) -> None:
"""
Should do nothing.
"""
before = """
x = "a"
"""
expected_after = """
x = "a"
"""
actual_after = self.insert_statements(
InsertAssignAroundIntegerVisitor(CodemodContext()), before
)
self.assertCodeEqual(expected_after, actual_after)
def test_metadata_works(self) -> None:
code = """
def foo() -> None:
pass
def bar() -> int:
return 5
"""
module = parse_module(dedent(code))
context = CodemodContext()
transform = TestingTransform(context)
transform.transform_module(module)
self.assertEqual(context.scratch, {
"foo": (2, 0),
"pass": (3, 4),
"bar": (5, 0)
})
def test_import_order(self) -> None:
"""
The imports should be in alphabetic order of added imports, added import alias, original imports.
"""
before = """
from a import b, e, h
"""
after = """
from a import c, f, d as x, g as y, b, e, h
"""
self.assertCodemod(
before,
after,
[("a", "f", None), ("a", "g", "y"), ("a", "c", None),
("a", "d", "x")],
context_override=CodemodContext(full_module_name="a.b.foobar"),
)
def _annotated_code(
stub: str,
code: str,
options: StubGenerationOptions,
) -> str:
"""
Merge inferred annotations from stubs with source code to get
annotated code.
"""
context = CodemodContext()
ApplyTypeAnnotationsVisitor.store_stub_in_context(
context=context,
stub=libcst.parse_module(stub),
use_future_annotations=options.use_future_annotations,
)
modified_tree = ApplyTypeAnnotationsVisitor(context).transform_module(
libcst.parse_module(code))
return modified_tree.code