diff --git a/DIRECTORY.md b/DIRECTORY.md index 0418ee0974f4..52b15110f236 100644 --- a/DIRECTORY.md +++ b/DIRECTORY.md @@ -867,6 +867,7 @@ * [Josephus Problem](maths/josephus_problem.py) * [Juggler Sequence](maths/juggler_sequence.py) * [Karatsuba](maths/karatsuba.py) + * [Keith Number](maths/keith_number.py) * [Kth Lexicographic Permutation](maths/kth_lexicographic_permutation.py) * [Laplace Transformation](maths/laplace_transformation.py) * [Largest Of Very Large Numbers](maths/largest_of_very_large_numbers.py) @@ -1492,6 +1493,7 @@ * [Non Preemptive Shortest Job First](scheduling/non_preemptive_shortest_job_first.py) * [Round Robin](scheduling/round_robin.py) * [Shortest Job First](scheduling/shortest_job_first.py) + * [Weighted Interval Scheduling](scheduling/weighted_interval_scheduling.py) ## [Searches](searches) * [Binary Search](searches/binary_search.py) @@ -1642,6 +1644,7 @@ * [Z Function](strings/z_function.py) ## [Tests](tests) + * [Test Find And Fix Pep798](tests/test_find_and_fix_pep798.py) * [Test Sorts](tests/test_sorts.py) ## [Web Programming](web_programming) diff --git a/scripts/find_and_fix_pep798.py b/scripts/find_and_fix_pep798.py new file mode 100755 index 000000000000..594d1deb8498 --- /dev/null +++ b/scripts/find_and_fix_pep798.py @@ -0,0 +1,245 @@ +#!/usr/bin/env -S uv run --script +# +# /// script +# requires-python = ">=3.13" +# dependencies = [] +# /// + +""" +Find (and optionally fix) opportunities to use PEP 798 in Python source files. + +PEP 798 allows unpacking inside comprehensions and generator expressions: +https://peps.python.org/pep-0798/ + +The following flattening patterns are detected: + + [y for x in xs for y in x] -> [*x for x in xs] + {y for x in xs for y in x} -> {*x for x in xs} + (y for x in xs for y in x) -> (*x for x in xs) + {k: v for d in ds for k, v in d.items()} -> {**d for d in ds} + +Only the simplest, always-equivalent forms are rewritten: exactly two ``for`` +clauses, no ``async``, and no ``if`` filter on the inner clause. Nested +candidates are fixed innermost first; re-run the script to catch any leftovers. +The rewritten source uses Python 3.15+ syntax, so fixed files need 3.15+. + +Usage: + scripts/find_and_fix_pep798.py [--fix] PATH [PATH ...] + +Without ``--fix`` the script only reports and exits with status 1 if any +opportunities were found. +""" + +import argparse +import ast +import sys +from collections.abc import Iterable, Iterator +from dataclasses import dataclass +from pathlib import Path + +Comprehension = ast.ListComp | ast.SetComp | ast.GeneratorExp | ast.DictComp + +# Nodes that need parentheses after ``*`` or ``**`` (they bind looser than ``|``). +_NEEDS_PARENS_AFTER_STAR = ( + ast.BoolOp, + ast.Compare, + ast.IfExp, + ast.Lambda, + ast.NamedExpr, + ast.UnaryOp, +) +# Nodes that need parentheses inside a ``for ... in`` / ``if`` comprehension clause. +_NEEDS_PARENS_IN_CLAUSE = (ast.IfExp, ast.Lambda, ast.NamedExpr, ast.Yield) + + +@dataclass(frozen=True) +class Finding: + """A PEP 798 opportunity: a source span and its replacement text.""" + + line: int + start: int + end: int + old: str + new: str + + +def _is_name(node: ast.expr, name: str) -> bool: + return isinstance(node, ast.Name) and node.id == name + + +def _is_flattening(node: Comprehension) -> bool: + """ + Return True if the comprehension is a simple two-level flattening. + + >>> def check(code: str) -> bool: + ... return _is_flattening(ast.parse(code, mode="eval").body) + >>> check("[y for x in xs for y in x]") + True + >>> check("[y for x in xs if x for y in x]") + True + >>> check("[y for x in xs for y in x if y]") + False + >>> check("[y + 1 for x in xs for y in x]") + False + >>> check("[y for x in xs]") + False + >>> check("{k: v for d in ds for k, v in d.items()}") + True + >>> check("{k: v for d in ds for k, v in d.items() if v}") + False + >>> check("{k: v for d in ds for v, k in d.items()}") + False + >>> check("{k: v for d in ds for k, v in d.items(1)}") + False + """ + if len(node.generators) != 2: + return False + outer, inner = node.generators + if outer.is_async or inner.is_async or inner.ifs: + return False + if isinstance(node, ast.DictComp): + match inner.target, node.key, node.value, inner.iter: + case ( + ast.Tuple(elts=[ast.Name(id=a), ast.Name(id=b)]), + ast.Name(id=key), + ast.Name(id=value), + ast.Call(func=ast.Attribute(attr="items"), args=[], keywords=[]), + ): + return a != b and (a, b) == (key, value) + return False + return isinstance(inner.target, ast.Name) and _is_name(node.elt, inner.target.id) + + +class _Offsets: + """Convert AST (line, UTF-8 byte column) positions to string offsets.""" + + def __init__(self, source: str) -> None: + self.lines = source.splitlines(keepends=True) + self.line_starts = [0] + for line in self.lines: + self.line_starts.append(self.line_starts[-1] + len(line)) + + def offset(self, line: int, byte_col: int) -> int: + text = self.lines[line - 1].encode()[:byte_col].decode() + return self.line_starts[line - 1] + len(text) + + +def _segment(source: str, offsets: _Offsets, node: ast.expr, parens: tuple) -> str: + start = offsets.offset(node.lineno, node.col_offset) + end = offsets.offset(node.end_lineno, node.end_col_offset) + text = source[start:end] + return f"({text})" if isinstance(node, parens) else text + + +def _replacement(source: str, offsets: _Offsets, node: Comprehension) -> str: + outer, inner = node.generators + star = "**" if isinstance(node, ast.DictComp) else "*" + iterable = inner.iter + if isinstance(node, ast.DictComp): + iterable = iterable.func.value + unpacked = _segment(source, offsets, iterable, _NEEDS_PARENS_AFTER_STAR) + target = _segment(source, offsets, outer.target, ()) + outer_iter = _segment(source, offsets, outer.iter, _NEEDS_PARENS_IN_CLAUSE) + clauses = [f"for {target} in {outer_iter}"] + clauses += [ + f"if {_segment(source, offsets, cond, _NEEDS_PARENS_IN_CLAUSE)}" + for cond in outer.ifs + ] + start = offsets.offset(node.lineno, node.col_offset) + end = offsets.offset(node.end_lineno, node.end_col_offset) + return f"{source[start]}{star}{unpacked} {' '.join(clauses)}{source[end - 1]}" + + +def _span(offsets: _Offsets, node: ast.AST) -> tuple[int, int]: + return ( + offsets.offset(node.lineno, node.col_offset), + offsets.offset(node.end_lineno, node.end_col_offset), + ) + + +def find_opportunities(source: str) -> list[Finding]: + """ + Return non-overlapping PEP 798 opportunities, innermost candidates first. + + >>> [f.new for f in find_opportunities("[y for x in xs for y in x]")] + ['[*x for x in xs]'] + >>> find_opportunities("[y for y in ys]") + [] + >>> len(find_opportunities("a = [y for x in xs for y in x]\\nb = {*x for x in xs}")) + 1 + """ + tree = ast.parse(source) + offsets = _Offsets(source) + candidates = [ + node + for node in ast.walk(tree) + if isinstance(node, Comprehension) and _is_flattening(node) + ] + # Innermost (smallest) first so that nested candidates keep the inner fix. + candidates.sort(key=lambda node: _span(offsets, node)[1] - _span(offsets, node)[0]) + taken: list[tuple[int, int]] = [] + findings: list[Finding] = [] + for node in candidates: + start, end = _span(offsets, node) + if any(start < t_end and t_start < end for t_start, t_end in taken): + continue + taken.append((start, end)) + findings.append( + Finding( + line=node.lineno, + start=start, + end=end, + old=source[start:end], + new=_replacement(source, offsets, node), + ) + ) + return sorted(findings, key=lambda finding: finding.start) + + +def fix_source(source: str) -> str: + """ + Rewrite every opportunity found by :func:`find_opportunities`. + + >>> print(fix_source("x = [y for ys in yss for y in ys]")) + x = [*ys for ys in yss] + >>> print(fix_source("x = {k: v for d in ds if d for k, v in d.items()}")) + x = {**d for d in ds if d} + """ + for finding in reversed(find_opportunities(source)): + source = f"{source[: finding.start]}{finding.new}{source[finding.end :]}" + return source + + +def iter_python_files(paths: Iterable[Path]) -> Iterator[Path]: + for path in paths: + if path.is_dir(): + yield from sorted(path.rglob("*.py")) + else: + yield path + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("paths", nargs="+", type=Path, help="files or directories") + parser.add_argument("--fix", action="store_true", help="rewrite files in place") + args = parser.parse_args(argv) + + total = 0 + for path in iter_python_files(args.paths): + source = path.read_text(encoding="utf-8") + try: + findings = find_opportunities(source) + except SyntaxError as error: + print(f"{path}: skipped ({error.msg})", file=sys.stderr) + continue + for finding in findings: + print(f"{path}:{finding.line}: {finding.old!r} -> {finding.new!r}") + total += len(findings) + if findings and args.fix: + path.write_text(fix_source(source), encoding="utf-8") + print(f"{total} PEP 798 opportunit{'y' if total == 1 else 'ies'} found") + return 0 if args.fix or not total else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_find_and_fix_pep798.py b/tests/test_find_and_fix_pep798.py new file mode 100644 index 000000000000..c74496347e58 --- /dev/null +++ b/tests/test_find_and_fix_pep798.py @@ -0,0 +1,146 @@ +from pathlib import Path + +import pytest + +from scripts.find_and_fix_pep798 import ( + find_opportunities, + fix_source, + iter_python_files, + main, +) + + +@pytest.mark.parametrize( + ("before", "after"), + [ + ("[y for x in xs for y in x]", "[*x for x in xs]"), + ("{y for x in xs for y in x}", "{*x for x in xs}"), + ("(y for x in xs for y in x)", "(*x for x in xs)"), + ("sum(y for x in xs for y in x)", "sum(*x for x in xs)"), + ("{k: v for d in ds for k, v in d.items()}", "{**d for d in ds}"), + ("[y for x in xs if x for y in x]", "[*x for x in xs if x]"), + ("[y for x in xs if x if x.a for y in x]", "[*x for x in xs if x if x.a]"), + ("[y for x, z in xs for y in x.items]", "[*x.items for x, z in xs]"), + ("[y for x in xs for y in x or ()]", "[*(x or ()) for x in xs]"), + ("[y for x in (a if b else c) for y in x]", "[*x for x in (a if b else c)]"), + ("[y for x in xs for y in (x)]", "[*x for x in xs]"), + ("[y for x in xs for y in f(x)[0]]", "[*f(x)[0] for x in xs]"), + ("[y for x in xs for y in x][0]", "[*x for x in xs][0]"), + ("[é for ü in xs for é in ü]", "[*ü for ü in xs]"), + ("{k: v for d in ds for k, v in (a or b).items()}", "{**(a or b) for d in ds}"), + ], +) +def test_fixes(before: str, after: str) -> None: + assert fix_source(before) == after + + +@pytest.mark.parametrize( + "source", + [ + "[y for y in ys]", + "[y + 1 for x in xs for y in x]", + "[z for x in xs for y in x]", + "[y for x in xs for y in x if y]", + "[y for x in xs for y in x for z in y]", + "[x for x in xs for x, y in x]", + "{k: v for d in ds for k, v in d.items() if v}", + "{k: v for d in ds for v, k in d.items()}", + "{k: v for d in ds for k, v in d.items(1)}", + "{k: v for d in ds for k, v in d}", + "{k: v for d in ds for k, k in d.items()}", + "[x for x in xs]", + "list(itertools.chain.from_iterable(xs))", + ], +) +def test_no_opportunity(source: str) -> None: + assert find_opportunities(source) == [] + assert fix_source(source) == source + + +def test_async_comprehensions_are_ignored() -> None: + source = ( + "async def f():\n" + " return [y async for x in xs for y in x]\n" + "async def g():\n" + " return [y for x in xs async for y in x]\n" + ) + assert find_opportunities(source) == [] + + +def test_find_reports_line_and_text() -> None: + source = "a = 1\nb = [y for x in xs for y in x]\n" + (finding,) = find_opportunities(source) + assert finding.line == 2 + assert finding.old == "[y for x in xs for y in x]" + assert finding.new == "[*x for x in xs]" + assert source[finding.start : finding.end] == finding.old + + +def test_multiple_findings_in_order() -> None: + source = ( + "a = [y for x in xs for y in x]\nb = {k: v for d in ds for k, v in d.items()}\n" + ) + assert fix_source(source) == "a = [*x for x in xs]\nb = {**d for d in ds}\n" + + +def test_nested_candidates_fix_innermost_only() -> None: + source = "[w for v in [y for x in xs for y in x] for w in v]" + (finding,) = find_opportunities(source) + assert finding.old == "[y for x in xs for y in x]" + assert fix_source(source) == "[w for v in [*x for x in xs] for w in v]" + + +def test_multiline_comprehension() -> None: + source = "result = [\n y\n for x in xs\n for y in x\n]\n" + assert fix_source(source) == "result = [*x for x in xs]\n" + + +def test_iter_python_files(tmp_path: Path) -> None: + (tmp_path / "sub").mkdir() + first = tmp_path / "a.py" + second = tmp_path / "sub" / "b.py" + other = tmp_path / "c.txt" + for path in (first, second, other): + path.write_text("") + assert list(iter_python_files([tmp_path])) == [first, second] + assert list(iter_python_files([other])) == [other] + + +def test_main_reports_without_fixing( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + path = tmp_path / "mod.py" + path.write_text("a = [y for x in xs for y in x]\n") + assert main([str(path)]) == 1 + output = capsys.readouterr().out + assert f"{path}:1:" in output + assert "1 PEP 798 opportunity found" in output + assert path.read_text() == "a = [y for x in xs for y in x]\n" + + +def test_main_fix_rewrites_file( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + path = tmp_path / "mod.py" + path.write_text("a = [y for x in xs for y in x]\nb = {y for x in xs for y in x}\n") + assert main(["--fix", str(path)]) == 0 + assert path.read_text() == "a = [*x for x in xs]\nb = {*x for x in xs}\n" + assert "2 PEP 798 opportunities found" in capsys.readouterr().out + + +def test_main_clean_file_returns_zero( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + path = tmp_path / "mod.py" + path.write_text("a = [y for y in ys]\n") + assert main([str(path)]) == 0 + assert "0 PEP 798 opportunities found" in capsys.readouterr().out + + +def test_main_skips_syntax_errors( + tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: + path = tmp_path / "bad.py" + path.write_text("def (:\n") + assert main([str(path)]) == 0 + assert "skipped" in capsys.readouterr().err