def test_findall_with_metadata_wrapper(self) -> None:
# Find all assignments in a tree
code = """
a = 1
b = True
def foo(bar: int) -> bool:
return False
"""
module = cst.parse_module(dedent(code))
wrapper = meta.MetadataWrapper(module)
# Test that when we find over a wrapper, we implicitly use it for
# metadata as well as traversal.
booleans = findall(
wrapper,
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.STORE),
)
self.assertNodeSequenceEqual(booleans, [cst.Name("a"), cst.Name("b")])
# Test that we can provide an explicit resolver and tree
booleans = findall(
wrapper.module,
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.STORE),
metadata_resolver=wrapper,
)
self.assertNodeSequenceEqual(booleans, [cst.Name("a"), cst.Name("b")])
# Test that failing to provide metadata leads to no match
booleans = findall(
wrapper.module,
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.STORE),
)
self.assertNodeSequenceEqual(booleans, [])
def test_remove_import_node(self) -> None:
"""
Make sure that if an import node itself is requested for
removal, we still do the right thing and only remove it
if it is unused.
"""
before = """
import foo
import qux
import bar as other
import foobar as other2
def fun() -> None:
qux.baz()
other2.baz()
"""
after = """
import qux
import foobar as other2
def fun() -> None:
qux.baz()
other2.baz()
"""
class RemoveImportTransformer(VisitorBasedCodemodCommand):
METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
def visit_Import(self, node: cst.Import) -> 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_remove_import_by_node_simple(self, before: str,
after: str) -> None:
"""
Given a node that's directly referenced in an import,
make sure that the import is removed when the node
is also removed.
"""
class RemoveBarTransformer(VisitorBasedCodemodCommand):
METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider)
@m.leave(
m.SimpleStatementLine(body=[
m.Expr(
m.Call(metadata=m.MatchMetadata(
QualifiedNameProvider,
{
QualifiedName(
source=QualifiedNameSource.IMPORT,
name="foo.bar",
)
},
)))
]))
def _leave_foo_bar(
self,
original_node: cst.SimpleStatementLine,
updated_node: cst.SimpleStatementLine,
) -> cst.RemovalSentinel:
RemoveImportsVisitor.remove_unused_import_by_node(
self.context, original_node)
return cst.RemoveFromParent()
module = cst.parse_module(self.make_fixture_data(before))
self.assertCodeEqual(
after,
RemoveBarTransformer(
CodemodContext()).transform_module(module).code)
def test_findall_with_transformers(self) -> None:
# Find all assignments in a tree
class TestTransformer(m.MatcherDecoratableTransformer):
METADATA_DEPENDENCIES: Sequence[meta.ProviderT] = (
meta.ExpressionContextProvider, )
def __init__(self) -> None:
super().__init__()
self.results: Sequence[cst.CSTNode] = ()
def visit_Module(self, node: cst.Module) -> None:
self.results = self.findall(
node,
m.MatchMetadata(meta.ExpressionContextProvider,
meta.ExpressionContext.STORE),
)
code = """
a = 1
b = True
def foo(bar: int) -> bool:
return False
"""
module = cst.parse_module(dedent(code))
wrapper = meta.MetadataWrapper(module)
visitor = TestTransformer()
wrapper.visit(visitor)
self.assertNodeSequenceEqual(
visitor.results,
[
cst.Name("a"),
cst.Name("b"),
cst.Name("foo"),
cst.Name("bar"),
],
)
def test_simple(self) -> None:
mock = Mock()
class ABatchable(BatchableCSTVisitor):
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_a()
object.__setattr__(node, "a_attr", True)
class BBatchable(BatchableCSTVisitor):
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_b()
object.__setattr__(node, "b_attr", 1)
module = visit_batched(parse_module("pass"), [ABatchable(), BBatchable()])
pass_ = cast(cst.SimpleStatementLine, module.body[0]).body[0]
# Check properties were set
self.assertEqual(object.__getattribute__(pass_, "a_attr"), True)
self.assertEqual(object.__getattribute__(pass_, "b_attr"), 1)
# Check that each visitor was only called once
mock.visited_a.assert_called_once()
mock.visited_b.assert_called_once()
def test_deep_remove_complex(self) -> None:
old_code = """
def a():
def b():
def c():
print("Hello, world!")
"""
new_code = """
def a():
def b():
pass
"""
module = cst.parse_module(dedent(old_code))
outer_fun = cst.ensure_type(module.body[0], cst.FunctionDef)
middle_fun = cst.ensure_type(
cst.ensure_type(outer_fun.body, cst.IndentedBlock).body[0],
cst.FunctionDef)
inner_fun = cst.ensure_type(
cst.ensure_type(middle_fun.body, cst.IndentedBlock).body[0],
cst.FunctionDef)
new_module = cst.ensure_type(module.deep_remove(inner_fun), cst.Module)
self.assertEqual(new_module.code, dedent(new_code))
def transform_module(self, tree: Module) -> Module:
"""
Transform entrypoint which handles multi-pass logic and metadata calculation
for you. This is the method that you should call if you wish to
invoke a codemod directly.
"""
if not self.should_allow_multiple_passes():
with self._handle_metadata_reference(tree) as tree_with_metadata:
return self.transform_module_impl(tree_with_metadata)
# We allow multiple passes, so we execute 1+ passes until there are
# no more changes.
before: str = tree.code
after: Optional[str] = None
while before != after:
if after is not None:
tree = parse_module(after)
before = after
with self._handle_metadata_reference(tree) as tree_with_metadata:
tree = self.transform_module_impl(tree_with_metadata)
after = tree.code
return tree
def _handle_string_annotation(
self, node: Union[cst.SimpleString, cst.ConcatenatedString]
) -> bool:
"""Returns whether it successfully handled the string annotation"""
if (
self.__in_type_hint_stack[-1] or self.__in_annotation_stack[-1]
) and not self.__in_ignored_subscript:
value = node.evaluated_value
if value:
top_level_annotation = self.__last_string_annotation is None
if top_level_annotation:
self.__last_string_annotation = node
try:
mod = cst.parse_module(value)
mod.visit(self)
except cst.ParserSyntaxError:
# swallow string annotation parsing errors
# this is the same behavior as cPython
pass
if top_level_annotation:
self.__last_string_annotation = None
return True
return False
def remove_target_typing_fields(self, files: List[Path]) -> None:
LOG.info("Removing typing options from %s targets files", len(files))
if self._pyre_only and not self._glob_threshold:
for path in files:
targets_file = Path(path)
source = targets_file.read_text()
output = libcst.parse_module(source).visit(
TargetPyreRemover()).code
targets_file.write_text(output)
else:
typing_options_regex = [
r"typing \?=.*",
r"check_types \?=.*",
r"check_types_options \?=.*",
r"typing_options \?=.*",
r"type_checker \?=.*",
]
remove_typing_fields_command = [
"sed",
"-i",
"/" + r"\|".join(typing_options_regex) + "/d",
] + [str(file) for file in files]
subprocess.run(remove_typing_fields_command)
def map_paths_operation(
path: Path,
rules: Set[LintRuleT],
type_cache: Optional[Mapping[ProviderT, object]],
) -> Union[str, Collection[BaseLintRuleReport]]:
# A top-level function to be accessible by `map_paths` from `fixit.cli`.
cst_wrapper = None
try:
if type_cache is not None:
cst_wrapper = MetadataWrapper(
cst.parse_module(SOURCE_CODE),
True,
type_cache,
)
return lint_file(
file_path=path,
source=SOURCE_CODE,
rules=rules,
cst_wrapper=cst_wrapper,
config=LintConfig(),
)
except Exception as e:
return str(e)
def test_has_name_helper(self) -> None:
class TestVisitor(cst.CSTVisitor):
METADATA_DEPENDENCIES = (QualifiedNameProvider, )
def __init__(self, test: UnitTest) -> None:
self.test = test
def visit_Call(self, node: cst.Call) -> Optional[bool]:
self.test.assertTrue(
QualifiedNameProvider.has_name(self, node, "a.b.c"))
self.test.assertFalse(
QualifiedNameProvider.has_name(self, node, "a.b"))
self.test.assertTrue(
QualifiedNameProvider.has_name(
self, node,
QualifiedName("a.b.c", QualifiedNameSource.IMPORT)))
self.test.assertFalse(
QualifiedNameProvider.has_name(
self, node,
QualifiedName("a.b.c", QualifiedNameSource.LOCAL)))
MetadataWrapper(cst.parse_module("import a;a.b.c()")).visit(
TestVisitor(self))
def test_gen_free_name():
src = '''
class P:
P5 = 1
def __init__(self): self.y = 0
def P0():
return P.P5
P1 = P0()
'''
tree = cst.parse_module(src)
env = SymbolTable({}, {})
free_name = gen_free_name(tree, env)
assert free_name == '_auto_name_0'
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P2'
env = SymbolTable({'P3': 'foo'}, {})
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P2'
env = SymbolTable({'P3': 'foo'}, {'P2': 'bar'})
free_name = gen_free_name(tree, env, prefix='P')
assert free_name == 'P4'
def test_replace_noop(self) -> None:
def _swap_bools(
node: cst.CSTNode,
extraction: Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]],
) -> cst.CSTNode:
return cst.Name("True" if cst.ensure_type(node, cst.Name).value ==
"False" else "False")
# Verify behavior when there's nothing to replace.
original = cst.parse_module(
"foo: int = 5\ndef bar() -> str:\n return 's'\n")
replaced = cst.ensure_type(
m.replace(
original,
m.Name("True") | m.Name("False"),
_swap_bools,
),
cst.Module,
)
# Should be identical tree contents
self.assertTrue(original.deep_equals(replaced))
# However, should be a new tree by identity
self.assertNotEqual(id(original), id(replaced))
def test_lambda_metadata_matcher(self) -> None:
# Match on qualified name provider
module = cst.parse_module(
"from typing import List\n\ndef foo() -> None: pass\n"
)
wrapper = cst.MetadataWrapper(module)
functiondef = cst.ensure_type(wrapper.module.body[1], cst.FunctionDef)
self.assertTrue(
matches(
functiondef,
m.FunctionDef(
name=m.MatchMetadataIfTrue(
meta.QualifiedNameProvider,
lambda qualnames: any(
n.name in {"foo", "bar", "baz"} for n in qualnames
),
)
),
metadata_resolver=wrapper,
)
)
self.assertFalse(
matches(
functiondef,
m.FunctionDef(
name=m.MatchMetadataIfTrue(
meta.QualifiedNameProvider,
lambda qualnames: any(
n.name in {"bar", "baz"} for n in qualnames
),
)
),
metadata_resolver=wrapper,
)
)
def exec_transform(src: str, transform: Transform) -> str:
py_ast = cst.parse_module(src)
old_ast = py_ast
if env.get('SIZR_DEBUG'):
print('<<<<< ORIGINAL:', py_ast)
transform_ctx = None
if transform.selector:
transform_ctx = select(py_ast, transform)
if env.get('SIZR_DEBUG'):
print('#> Matches <#########################')
for m in transform_ctx.matches:
print(m)
print('#####################################')
if transform.assertion:
py_ast = assert_(py_ast, transform_ctx)
if env.get('SIZR_DEBUG') and transform.assertion:
print('>>>>> TRANSFORMED:', py_ast)
print('!!!!! DIFF', ''.join(
difflib.unified_diff(
str(old_ast).splitlines(1),
str(py_ast).splitlines(1)
)
))
result = py_ast.code
diff = ''.join(
difflib.unified_diff(
src.splitlines(1),
result.splitlines(1)
)
)
if diff:
print(diff)
else:
print('no changes!')
return result
def test_batched_provider(self) -> None:
"""
Tests that batchable providers are resolved correctly.
Sets metadata on:
- pass: BatchedProviderA -> 1
BatchedProviderB -> "a"
"""
test = self
mock = Mock()
class BatchedProviderA(BatchableMetadataProvider[int]):
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_a()
self.set_metadata(node, 1)
class BatchedProviderB(BatchableMetadataProvider[str]):
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_b()
self.set_metadata(node, "a")
class DependentVisitor(CSTTransformer):
METADATA_DEPENDENCIES = (BatchedProviderA, BatchedProviderB)
def visit_Pass(self, node: cst.Pass) -> None:
# Check metadata is set
test.assertEqual(self.get_metadata(BatchedProviderA, node), 1)
test.assertEqual(self.get_metadata(BatchedProviderB, node),
"a")
module = parse_module("pass")
module.visit(DependentVisitor())
# Check that each batchable visitor is only called once
mock.visited_a.assert_called_once()
mock.visited_b.assert_called_once()
def _is_awaitable_callable(annotation: str) -> bool:
if not annotation.startswith("typing.Callable"):
# Exit early if this is not even a `typing.Callable` annotation.
return False
try:
# Wrap this in a try-except since the type annotation may not be parse-able as a module.
# If it is not parse-able, we know it's not what we are looking for anyway, so return `False`.
parsed_ann = cst.parse_module(annotation)
except Exception:
return False
# If passed annotation does not match the expected annotation structure for a `typing.Callable` with
# typing.Coroutine as the return type, matched_callable_ann will simply be `None`.
# The expected structure of an awaitable callable annotation from Pyre is: typing.Callable()[[...], typing.Coroutine[...]]
matched_callable_ann: Optional[Dict[str, Union[
Sequence[cst.CSTNode], cst.CSTNode]]] = m.extract(
parsed_ann,
m.Module(body=[
m.SimpleStatementLine(body=[
m.Expr(value=m.Subscript(slice=[
m.SubscriptElement(),
m.SubscriptElement(slice=m.Index(value=m.Subscript(
value=m.SaveMatchedNode(
m.Attribute(),
"base_return_type",
)))),
], ))
]),
]),
)
if (matched_callable_ann is not None
and "base_return_type" in matched_callable_ann):
base_return_type = get_full_name_for_node(
cst.ensure_type(matched_callable_ann["base_return_type"],
cst.CSTNode))
return (base_return_type is not None
and base_return_type == "typing.Coroutine")
return False
def parse(model_fn: FunctionType) -> Tuple[GraphicalModel, dict]:
"""Parse the model definition to build a graphical model.
Parameter
---------
model
A live function object that contains the model definition.
Returns
-------
The intermediate representation of the model.
"""
source = inspect.getsource(model_fn)
source = textwrap.dedent(source) # TODO: do we need this with libcst?
tree = cst.parse_module(source)
namespace = model_fn.__globals__
definition_visitor = ModelDefinitionParser(namespace)
tree.visit(definition_visitor)
graph = definition_visitor.graph
return graph, namespace
def test_all_visits(self) -> None:
mock = Mock()
class Batchable(BatchableCSTVisitor):
def visit_Pass(self, node: cst.Pass) -> None:
mock.visit_Pass()
object.__setattr__(node, "visit_Pass", True)
def visit_Pass_semicolon(self, node: cst.Pass) -> None:
mock.visit_Pass_semicolon()
object.__setattr__(node, "visit_Pass_semicolon", True)
def leave_Pass_semicolon(self, node: cst.Pass) -> None:
mock.leave_Pass_semicolon()
object.__setattr__(node, "leave_Pass_semicolon", True)
def leave_Pass(self, original_node: cst.Pass) -> None:
mock.leave_Pass()
object.__setattr__(original_node, "leave_Pass", True)
module = visit_batched(parse_module("pass"), [Batchable()])
pass_ = cast(cst.SimpleStatementLine, module.body[0]).body[0]
# Check properties were set
self.assertEqual(object.__getattribute__(pass_, "visit_Pass"), True)
self.assertEqual(object.__getattribute__(pass_, "leave_Pass"), True)
self.assertEqual(
object.__getattribute__(pass_, "visit_Pass_semicolon"), True)
self.assertEqual(
object.__getattribute__(pass_, "leave_Pass_semicolon"), True)
# Check that each visitor was only called once
mock.visit_Pass.assert_called_once()
mock.leave_Pass.assert_called_once()
mock.visit_Pass_semicolon.assert_called_once()
mock.leave_Pass_semicolon.assert_called_once()
def _line_ranges_spanned_by_format_strings(
source: str, ) -> Dict[libcst.CSTNode, LineRange]:
def _code_range_to_line_range(
code_range: libcst._position.CodeRange, ) -> LineRange:
return code_range.start.line, code_range.end.line
try:
wrapper = libcst.metadata.MetadataWrapper(libcst.parse_module(source))
except libcst._exceptions.ParserSyntaxError as exception:
# NOTE: This should not happen. If a file is unparseable for libcst, it
# would probably have been unparseable for Pyre as well. In that case,
# we would not have raised a 404 parse error and not reached here in the
# first place. Still, catch the exception and just skip the special
# handling of format strings.
LOG.warning("Not moving out fixmes from f-strings because"
f" libcst failed to parse the file: {exception}")
return {}
position_map = wrapper.resolve(libcst.metadata.PositionProvider)
return {
format_string: _code_range_to_line_range(position_map[format_string])
for format_string in libcst_matchers.findall(
wrapper.module, libcst_matchers.FormattedString())
}
def test_simple(self) -> None:
mock = Mock()
class ABatchable(BatchableCSTVisitor):
def visit_Del(self, node: cst.Del) -> None:
object.__setattr__(node, "target", mock.visited_a())
class BBatchable(BatchableCSTVisitor):
def visit_Del(self, node: cst.Del) -> None:
object.__setattr__(node, "semicolon", mock.visited_b())
module = visit_batched(parse_module("del a"),
[ABatchable(), BBatchable()])
del_ = cast(cst.SimpleStatementLine, module.body[0]).body[0]
# Check that each visitor was only called once
mock.visited_a.assert_called_once()
mock.visited_b.assert_called_once()
# Check properties were set
self.assertEqual(object.__getattribute__(del_, "target"),
mock.visited_a())
self.assertEqual(object.__getattribute__(del_, "semicolon"),
mock.visited_b())
def test_copies_tree(self) -> None:
m = cst.parse_module("pass")
mw = MetadataWrapper(m)
self.assertTrue(mw.module.deep_equals(m))
self.assertIsNot(mw.module, m)
def test_unsafe_skip_copy(self) -> None:
m = cst.parse_module("pass")
mw = MetadataWrapper(m, unsafe_skip_copy=True)
self.assertIs(mw.module, m)
def _parse(file: IO[str]) -> libcst.Module:
contents = file.read()
return libcst.parse_module(contents)
def test_annotate_functions_py38(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 format_files(source: str) -> Module:
return parse_module(textwrap.dedent(source.rstrip()))
def test_mixed_providers(self) -> None:
"""
Tests that a mixed set of providers is resolved properly.
Sets metadata on pass:
BatchedProviderA -> 2
BatchedProviderB -> 3
DependentProvider -> 5
DependentBatched -> 4
"""
test = self
mock = Mock()
class SimpleProvider(VisitorMetadataProvider[int]):
def visit_Pass(self, node: cst.CSTNode) -> None:
mock.visited_simple()
self.set_metadata(node, 1)
class BatchedProviderA(BatchableMetadataProvider[int]):
METADATA_DEPENDENCIES = (SimpleProvider, )
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_a()
self.set_metadata(node, 2)
class BatchedProviderB(BatchableMetadataProvider[int]):
METADATA_DEPENDENCIES = (SimpleProvider, )
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_b()
self.set_metadata(node, 3)
class DependentProvider(VisitorMetadataProvider[int]):
METADATA_DEPENDENCIES = (BatchedProviderA, BatchedProviderB)
def on_visit(self, node: cst.CSTNode) -> bool:
sum = self.get_metadata(BatchedProviderA, node,
0) + self.get_metadata(
BatchedProviderB, node, 0)
self.set_metadata(node, sum)
return True
class BatchedProviderC(BatchableMetadataProvider[int]):
METADATA_DEPENDENCIES = (BatchedProviderA, )
def visit_Pass(self, node: cst.Pass) -> None:
mock.visited_c()
self.set_metadata(
node,
self.get_metadata(BatchedProviderA, node) * 2)
class DependentVisitor(CSTTransformer):
METADATA_DEPENDENCIES = (
BatchedProviderA,
BatchedProviderB,
BatchedProviderC,
DependentProvider,
)
def visit_Module(self, module: cst.Module) -> None:
# Dependent visitor set metadata on all nodes but for module it
# defaulted to 0 because BatchedProviderA/B only set metadata on
# pass nodes
test.assertEqual(self.get_metadata(DependentProvider, module),
0)
def visit_Pass(self, node: cst.Pass) -> None:
# Check metadata is set
test.assertEqual(self.get_metadata(BatchedProviderA, node), 2)
test.assertEqual(self.get_metadata(BatchedProviderB, node), 3)
test.assertEqual(self.get_metadata(BatchedProviderC, node), 4)
test.assertEqual(self.get_metadata(DependentProvider, node), 5)
module = parse_module("pass")
module.visit(DependentVisitor())
# Check each visitor is called once
mock.visited_simple.assert_called_once()
mock.visited_a.assert_called_once()
mock.visited_b.assert_called_once()
mock.visited_c.assert_called_once()
# https://libcst.readthedocs.io/en/latest/_modules/libcst/codemod/visitors/_apply_type_annotations.html#ApplyTypeAnnotationsVisitor
# # Only add new annotation if explicitly told to overwrite existing
# # annotations or if one doesn't already exist.
# if self.overwrite_existing_annotations or not updated_node.returns:
# updated_node = updated_node.with_changes(
# returns=function_annotation.returns
# )
# # Don't override default values when annotating functions
# new_parameters = self._update_parameters(function_annotation, updated_node)
# return updated_node.with_changes(params=new_parameters)
# %%
# %%
source_tree = cst.parse_module(simple_stub)
info_tree = cst.parse_module(rich_source)
# %%
visitor = TypingCollector()
info_tree.visit(visitor)
# %%
transformer = TypingTransformer(visitor.annotations)
modified_tree = source_tree.visit(transformer)
# %%
print("=" * 20)
print(modified_tree.code)
# %%
def parallel_exec_transform_with_prettyprint( # noqa: C901
transform: Codemod,
files: Sequence[str],
*,
jobs: Optional[int] = None,
unified_diff: Optional[int] = None,
include_generated: bool = False,
generated_code_marker: str = _DEFAULT_GENERATED_CODE_MARKER,
format_code: bool = False,
formatter_args: Sequence[str] = (),
show_successes: bool = False,
hide_generated: bool = False,
hide_blacklisted: bool = False,
hide_progress: bool = False,
blacklist_patterns: Sequence[str] = (),
python_version: Optional[str] = None,
repo_root: Optional[str] = None,
) -> ParallelTransformResult:
"""
Given a list of files and an instantiated codemod we should apply to them,
fork and apply the codemod in parallel to all of the files, including any
configured formatter. The ``jobs`` parameter controls the maximum number of
in-flight transforms, and needs to be at least 1. If not included, the number
of jobs will automatically be set to the number of CPU cores. If ``unified_diff``
is set to a number, changes to files will be printed to stdout with
``unified_diff`` lines of context. If it is set to ``None`` or left out, files
themselves will be updated with changes and formatting. If a
``python_version`` is provided, then we will parse each source file using
this version. Otherwise, we will use the version of the currently executing python
binary.
A progress indicator as well as any generated warnings will be printed to stderr.
To supress the interactive progress indicator, set ``hide_progress`` to ``True``.
Files that include the generated code marker will be skipped unless the
``include_generated`` parameter is set to ``True``. Similarly, files that match
a supplied blacklist of regex patterns will be skipped. Warnings for skipping
both blacklisted and generated files will be printed to stderr along with
warnings generated by the codemod unless ``hide_blacklisted`` and
``hide_generated`` are set to ``True``. Files that were successfully codemodded
will not be printed to stderr unless ``show_successes`` is set to ``True``.
To make this API possible, we take an instantiated transform. This is due to
the fact that lambdas are not pickleable and pickling functions is undefined.
This means we're implicitly relying on fork behavior on UNIX-like systems, and
this function will not work on Windows systems. To create a command-line utility
that runs on Windows, please instead see
:func:`~libcst.codemod.exec_transform_with_prettyprint`.
"""
# Ensure that we have no duplicates, otherwise we might get race conditions
# on write.
files = sorted(list({os.path.abspath(f) for f in files}))
total = len(files)
progress = Progress(enabled=not hide_progress, total=total)
# Grab number of cores if we need to
jobs: int = jobs if jobs is not None else cpu_count()
if jobs < 1:
raise Exception("Must have at least one job to process!")
if total == 0:
return ParallelTransformResult(successes=0,
failures=0,
skips=0,
warnings=0)
if repo_root is not None:
# Make sure if there is a root that we have the absolute path to it.
repo_root = os.path.abspath(repo_root)
# Spin up a full repo metadata manager so that we can provide metadata
# like type inference to individual forked processes.
print("Calculating full-repo metadata...", file=sys.stderr)
metadata_manager = FullRepoManager(
repo_root,
files,
transform.get_inherited_dependencies(),
)
metadata_manager.resolve_cache()
transform.context = replace(
transform.context,
metadata_manager=metadata_manager,
)
print("Executing codemod...", file=sys.stderr)
config = ExecutionConfig(
repo_root=repo_root,
unified_diff=unified_diff,
include_generated=include_generated,
generated_code_marker=generated_code_marker,
format_code=format_code,
formatter_args=formatter_args,
blacklist_patterns=blacklist_patterns,
python_version=python_version,
)
if total == 1:
# Simple case, we should not pay for process overhead.
# Let's just use a dummy synchronous pool.
jobs = 1
pool_impl = DummyPool
else:
pool_impl = Pool
# Warm the parser, pre-fork.
parse_module(
"",
config=(PartialParserConfig(python_version=python_version)
if python_version is not None else PartialParserConfig()),
)
successes: int = 0
failures: int = 0
warnings: int = 0
skips: int = 0
with pool_impl(processes=jobs) as p: # type: ignore
args = [{
"transformer": transform,
"filename": filename,
"config": config,
} for filename in files]
try:
for result in p.imap_unordered(_execute_transform_wrap,
args,
chunksize=4):
# Print an execution result, keep track of failures
_print_parallel_result(
result,
progress,
unified_diff=bool(unified_diff),
show_successes=show_successes,
hide_generated=hide_generated,
hide_blacklisted=hide_blacklisted,
)
progress.print(successes + failures + skips)
if isinstance(result.transform_result, TransformFailure):
failures += 1
elif isinstance(result.transform_result, TransformSuccess):
successes += 1
elif isinstance(result.transform_result,
(TransformExit, TransformSkip)):
skips += 1
warnings += len(result.transform_result.warning_messages)
finally:
progress.clear()
# Return whether there was one or more failure.
return ParallelTransformResult(successes=successes,
failures=failures,
skips=skips,
warnings=warnings)
def _execute_transform( # noqa: C901
transformer: Codemod,
filename: str,
config: ExecutionConfig,
) -> ExecutionResult:
for pattern in config.blacklist_patterns:
if re.fullmatch(pattern, filename):
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformSkip(
skip_reason=SkipReason.BLACKLISTED,
skip_description=f"Blacklisted by pattern {pattern}.",
),
)
try:
with open(filename, "rb") as fp:
oldcode = fp.read()
# Skip generated files
if (not config.include_generated
and config.generated_code_marker.encode("utf-8") in oldcode):
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformSkip(
skip_reason=SkipReason.GENERATED,
skip_description="Generated file.",
),
)
# Somewhat gross hack to provide the filename in the transform's context.
# We do this after the fork so that a context that was initialized with
# some defaults before calling parallel_exec_transform_with_prettyprint
# will be updated per-file.
transformer.context = replace(
transformer.context,
filename=filename,
full_module_name=_calculate_module(config.repo_root, filename),
)
# Run the transform, bail if we failed or if we aren't formatting code
try:
input_tree = parse_module(
oldcode,
config=(PartialParserConfig(
python_version=str(config.python_version))
if config.python_version is not None else
PartialParserConfig()),
)
output_tree = transformer.transform_module(input_tree)
newcode = output_tree.bytes
encoding = output_tree.encoding
except KeyboardInterrupt:
return ExecutionResult(filename=filename,
changed=False,
transform_result=TransformExit())
except SkipFile as ex:
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformSkip(
skip_reason=SkipReason.OTHER,
skip_description=str(ex),
warning_messages=transformer.context.warnings,
),
)
except Exception as ex:
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformFailure(
error=ex,
traceback_str=traceback.format_exc(),
warning_messages=transformer.context.warnings,
),
)
# Call formatter if needed, but only if we actually changed something in this
# file
if config.format_code and newcode != oldcode:
try:
newcode = invoke_formatter(config.formatter_args, newcode)
except KeyboardInterrupt:
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformExit(),
)
except Exception as ex:
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformFailure(
error=ex,
traceback_str=traceback.format_exc(),
warning_messages=transformer.context.warnings,
),
)
# Format as unified diff if needed, otherwise save it back
changed = oldcode != newcode
if config.unified_diff:
newcode = diff_code(
oldcode.decode(encoding),
newcode.decode(encoding),
config.unified_diff,
filename=filename,
)
else:
# Write back if we changed
if changed:
with open(filename, "wb") as fp:
fp.write(newcode)
# Not strictly necessary, but saves space in pickle since we won't use it
newcode = ""
# Inform success
return ExecutionResult(
filename=filename,
changed=changed,
transform_result=TransformSuccess(
warning_messages=transformer.context.warnings, code=newcode),
)
except KeyboardInterrupt:
return ExecutionResult(filename=filename,
changed=False,
transform_result=TransformExit())
except Exception as ex:
return ExecutionResult(
filename=filename,
changed=False,
transform_result=TransformFailure(
error=ex,
traceback_str=traceback.format_exc(),
warning_messages=transformer.context.warnings,
),
)
