1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
|
import contextlib
import os
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Dict, Generator
from unittest import TestCase
from libcst import BaseExpression, Call, matchers as m, Name
from libcst.codemod import (
CodemodContext,
gather_files,
parallel_exec_transform_with_prettyprint,
VisitorBasedCodemodCommand,
)
from libcst.codemod.visitors import AddImportsVisitor
class PrintToPPrintCommand(VisitorBasedCodemodCommand):
def __init__(self, context: CodemodContext, **kwargs: Dict[str, object]) -> None:
super().__init__(context, **kwargs)
self.context.scratch["PPRINT_WAS_HERE"] = True
def leave_Call(self, original_node: Call, updated_node: Call) -> BaseExpression:
if not self.context.scratch["PPRINT_WAS_HERE"]:
raise AssertionError("Scratch space lost")
if m.matches(updated_node, m.Call(func=m.Name("print"))):
AddImportsVisitor.add_needed_import(
self.context,
"pprint",
"pprint",
)
return updated_node.with_changes(func=Name("pprint"))
return super().leave_Call(original_node, updated_node)
@contextlib.contextmanager
def temp_workspace() -> Generator[Path, None, None]:
cwd = os.getcwd()
with TemporaryDirectory() as temp_dir:
try:
ws = Path(temp_dir).resolve()
os.chdir(ws)
yield ws
finally:
os.chdir(cwd)
class ToolE2ETest(TestCase):
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()
# Just a dir named "dir.py", should be ignored
adir = tmp / "dir.py"
adir.mkdir()
# Run command
command_instance = PrintToPPrintCommand(CodemodContext())
files = gather_files(".")
result = parallel_exec_transform_with_prettyprint(
command_instance,
files,
format_code=False,
hide_progress=True,
)
print(result)
# 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",
)
|
