def shed(
*, source_code: str,
first_party_imports: FrozenSet[str] = frozenset()) -> str:
"""Process the source code of a single module."""
assert isinstance(source_code, str)
assert isinstance(first_party_imports, frozenset)
assert all(isinstance(name, str) for name in first_party_imports)
assert all(name.isidentifier() for name in first_party_imports)
# Use black to autodetect our target versions
target_versions = {
v
for v in black.detect_target_versions(
black.lib2to3_parse(source_code.lstrip(), set(_version_map)))
if v.value >= black.TargetVersion.PY36.value
}
assert target_versions
input_code = source_code
# Autoflake first:
source_code = autoflake.fix_code(
source_code,
expand_star_imports=True,
remove_all_unused_imports=True,
remove_duplicate_keys=True,
remove_unused_variables=True,
)
# Then isort...
# TODO: swap as soon as 5.0 is released for black compat & clean config handling
# source_code = isort.api.sorted_imports(
# file_contents=source_code, known_first_party=first_party_imports,
# )
source_code = isort.SortImports(file_contents=source_code).output
# Now pyupgrade - see pyupgrade._fix_file
source_code = pyupgrade._fix_tokens(
source_code,
min_version=_version_map[min(target_versions,
key=attrgetter("value"))],
)
source_code = pyupgrade._fix_percent_format(source_code)
source_code = pyupgrade._fix_py3_plus(source_code)
# and finally Black!
source_code = black.format_str(
source_code, mode=black.FileMode(target_versions=target_versions))
if source_code == input_code:
return source_code
# If we've modified the code, iterate to a fixpoint.
# e.g. "pass;#" -> "pass\n#\n" -> "#\n"
return shed(source_code=source_code,
first_party_imports=first_party_imports)
def test_future_remove(s, min_version, expected):
assert _fix_tokens(s, min_version=min_version) == expected
def test_noop(s, min_version):
assert _fix_tokens(s, min_version=min_version) == s
def test_unicode_literals_noop(s, min_version):
assert _fix_tokens(s, min_version=min_version) == s
def test_binary_literals(s, expected):
assert _fix_tokens(s, min_version=(2, 7)) == expected
def test_fix_octal_literal(s, expected):
assert _fix_tokens(s) == expected
def test_long_literals(s, expected):
assert _fix_tokens(s) == expected
def test_fix_octal_literal(s, expected):
assert _fix_tokens(s, py3_plus=False) == expected
def test_fix_escape_sequences_noop(s):
assert _fix_tokens(s, py3_plus=False) == s
def test_fix_extra_parens(s, expected):
assert _fix_tokens(s, py3_plus=False) == expected
def test_fix_extra_parens_noop(s):
assert _fix_tokens(s, py3_plus=False) == s
def test_unicode_literals(s, py3_plus, expected):
ret = _fix_tokens(s, py3_plus=py3_plus)
assert ret == expected
def test_fix_extra_parens(s, expected):
assert _fix_tokens(s, min_version=(2, 7)) == expected
def test_fix_extra_parens_noop(s):
assert _fix_tokens(s, min_version=(2, 7)) == s
def test_future_remove_noop(s, min_version):
assert _fix_tokens(s, min_version=min_version) == s
def test_binary_literals(s, expected):
assert _fix_tokens(s, py3_plus=False) == expected
def test_noop_octal_literals(s):
assert _fix_tokens(s, py3_plus=False) == s
def test_fix_escape_sequences(s, expected):
assert _fix_tokens(s, py3_plus=False) == expected
def test_fix_ur_literals_gets_fixed_before_u_removed():
assert _fix_tokens("ur'\\s\\u2603'", min_version=(3, )) == "'\\\\s\\u2603'"
def test_fix_ur_literals(s, expected):
ret = _fix_tokens(s, py3_plus=False)
assert ret == expected
def test_noop_octal_literals(s):
assert _fix_tokens(s) == s
def test_fix_ur_literals_gets_fixed_before_u_removed():
assert _fix_tokens("ur'\\s\\u2603'", py3_plus=True) == "'\\\\s\\u2603'"
def test_binary_literals_noop(s):
assert _fix_tokens(s, min_version=(2, 7)) == s
def test_import_removals(s, min_version, expected):
assert _fix_tokens(s, min_version=min_version) == expected
def shed(
source_code: str,
*,
refactor: bool = False,
first_party_imports: FrozenSet[str] = frozenset(),
) -> str:
"""Process the source code of a single module."""
assert isinstance(source_code, str)
assert isinstance(refactor, bool)
assert isinstance(first_party_imports, frozenset)
assert all(isinstance(name, str) for name in first_party_imports)
assert all(name.isidentifier() for name in first_party_imports)
# Use black to autodetect our target versions
target_versions = {
v
for v in black.detect_target_versions(
black.lib2to3_parse(source_code.lstrip(), set(_version_map)))
if v.value >= black.TargetVersion.PY36.value
}
assert target_versions
min_version = _version_map[min(target_versions, key=attrgetter("value"))]
if refactor:
# Some tools assume that the file is multi-line, but empty files are valid input.
source_code += "\n"
# Use com2ann to comvert type comments to annotations on Python 3.8+
source_code, _ = com2ann(
source_code,
drop_ellipsis=True,
silent=True,
python_minor_version=min_version[1],
)
# Use teyit to replace old unittest.assertX methods on Python 3.9+
source_code, _ = _teyit_refactor(source_code)
# Then apply pybetter's fixes with libcst
tree = libcst.parse_module(source_code)
for fixer in _pybetter_fixers:
tree = fixer(tree)
source_code = tree.code
# Then shed.docshed (below) formats any code blocks in documentation
source_code = docshed(source=source_code,
first_party_imports=first_party_imports)
# And pyupgrade - see pyupgrade._fix_file - is our last stable fixer
# Calculate separate minver because pyupgrade doesn't have py39-specific logic yet
pyupgrade_min_ver = min(min_version, max(pyupgrade.IMPORT_REMOVALS.keys()))
source_code = pyupgrade._fix_tokens(source_code,
min_version=pyupgrade_min_ver)
source_code = pyupgrade._fix_percent_format(source_code)
source_code = pyupgrade._fix_py3_plus(source_code,
min_version=pyupgrade_min_ver)
source_code = pyupgrade._fix_py36_plus(source_code)
# One tricky thing: running `isort` or `autoflake` can "unlock" further fixes
# for `black`, e.g. "pass;#" -> "pass\n#\n" -> "#\n". We therefore loop until
# neither of them have made a change in the last loop body, trusting that
# `black` itself is idempotent because that's tested upstream.
prev = ""
black_mode = black.FileMode(target_versions=target_versions)
while prev != source_code:
prev = source_code = black.format_str(source_code, mode=black_mode)
source_code = autoflake.fix_code(
source_code,
expand_star_imports=True,
remove_all_unused_imports=True,
remove_duplicate_keys=True,
remove_unused_variables=True,
)
source_code = isort.code(
source_code,
known_first_party=first_party_imports,
profile="black",
combine_as_imports=True,
)
# Remove any extra trailing whitespace
return source_code.rstrip() + "\n"
def test_import_removals_noop(s, min_version):
assert _fix_tokens(s, min_version=min_version) == s
def test_unicode_literals(s, min_version, expected):
assert _fix_tokens(s, min_version=min_version) == expected
def test_binary_literals_noop(s):
assert _fix_tokens(s, py3_plus=False) == s
def test_rewrite(s, expected):
assert _fix_tokens(s, min_version=(3, )) == expected
def test_unicode_literals_noop(s, py3_plus):
assert _fix_tokens(s, py3_plus=py3_plus) == s
