Add scripts/find_and_fix_pep798.py

@@ -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)

@@ -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())

@@ -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