misc.python.materialize.cli.fmt

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10"""fmt — formats Rust, Python & Protobuf files in parallel."""
 11
 12import argparse
 13import json
 14import math
 15import os
 16import subprocess
 17
 18from materialize import MZ_ROOT
 19from materialize.parallel_task import TaskSpec, run_parallel
 20
 21# Rust sources live in more than one workspace, and `cargo metadata` only ever
 22# reports the one it is pointed at. The `src/*/fuzz` cargo-fuzz crates attach to
 23# the `test/cargo-fuzz` workspace, which the root workspace does not include, so
 24# without a second invocation here they go unformatted entirely.
 25RUST_MANIFESTS = ["Cargo.toml", "test/cargo-fuzz/Cargo.toml"]
 26
 27
 28def main() -> int:
 29    parser = argparse.ArgumentParser(prog="fmt")
 30    parser.add_argument("--check", action="store_true")
 31    args = parser.parse_args()
 32
 33    tasks: list[tuple[str, TaskSpec]] = [
 34        ("rustfmt", _rustfmt_fn(check=args.check)),
 35        ("buf", _buf_cmd(check=args.check)),
 36    ]
 37
 38    if not os.environ.get("MZDEV_NO_PYTHON"):
 39        tasks += [
 40            ("black", _black_cmd(check=args.check)),
 41            ("ruff", _ruff_cmd(check=args.check)),
 42            ("ruff-dbt", _ruff_dbt_cmd(check=args.check)),
 43        ]
 44
 45    return 1 if run_parallel(tasks, spinner_suffix="formatters") else 0
 46
 47
 48def _rustfmt_fn(*, check: bool):
 49    """Return a callable that runs cargo metadata + parallel rustfmt."""
 50
 51    def run() -> tuple[bool, str]:
 52        ncpus = os.cpu_count() or 8
 53
 54        kinds = {"lib", "bin", "bench", "test", "example", "proc-macro", "custom-build"}
 55        # Keyed by edition: `gen` is an identifier in 2021 but a reserved keyword
 56        # in 2024, so rustfmt cannot even parse a file at the wrong edition.
 57        paths_by_edition: dict[str, list[str]] = {}
 58        for manifest in RUST_MANIFESTS:
 59            result = subprocess.run(
 60                [
 61                    "cargo",
 62                    "metadata",
 63                    "--no-deps",
 64                    "--format-version=1",
 65                    f"--manifest-path={manifest}",
 66                ],
 67                capture_output=True,
 68                text=True,
 69            )
 70            if result.returncode != 0:
 71                return False, result.stderr.strip()
 72
 73            meta = json.loads(result.stdout)
 74            for pkg in meta["packages"]:
 75                for t in pkg["targets"]:
 76                    if kinds & set(t["kind"]):
 77                        paths_by_edition.setdefault(pkg["edition"], []).append(
 78                            t["src_path"]
 79                        )
 80        if not paths_by_edition:
 81            return True, ""
 82
 83        # Split into batches and run rustfmt in parallel.
 84        batches = []
 85        for edition, paths in paths_by_edition.items():
 86            batch_size = math.ceil(len(paths) / ncpus)
 87            batches += [
 88                (edition, paths[i : i + batch_size])
 89                for i in range(0, len(paths), batch_size)
 90            ]
 91
 92        cmd_base = ["rustfmt", "--config", "error_on_line_overflow=true"]
 93        if check:
 94            cmd_base.append("--check")
 95
 96        procs = [
 97            subprocess.Popen(
 98                cmd_base + [f"--edition={edition}"] + batch,
 99                stdout=subprocess.PIPE,
100                stderr=subprocess.STDOUT,
101            )
102            for edition, batch in batches
103        ]
104
105        all_output = []
106        all_ok = True
107        for proc in procs:
108            stdout, _ = proc.communicate()
109            if proc.returncode != 0:
110                all_ok = False
111            out = stdout.decode("utf-8").strip()
112            if out:
113                all_output.append(out)
114
115        return all_ok, "\n".join(all_output)
116
117    return run
118
119
120def _buf_cmd(*, check: bool) -> list[str]:
121    if check:
122        return ["buf", "format", "src", "--diff", "--exit-code"]
123    return ["buf", "format", "src", "-w"]
124
125
126def _black_cmd(*, check: bool) -> list[str]:
127    args = "--check --quiet" if check else "--quiet"
128    return [
129        "bash",
130        "-c",
131        f'. misc/shlib/shlib.bash && git_files "*.py" | xargs bin/pyactivate -m black {args}',
132    ]
133
134
135def _ruff_cmd(*, check: bool) -> list[str]:
136    fix = "" if check else " --fix"
137    return [
138        "bash",
139        "-c",
140        f'. misc/shlib/shlib.bash && git_files "*.py" | grep -v "^misc/dbt-materialize/" | xargs bin/pyactivate -m ruff{fix}',
141    ]
142
143
144def _ruff_dbt_cmd(*, check: bool) -> list[str]:
145    fix = "" if check else " --fix"
146    return [
147        "bash",
148        "-c",
149        f'. misc/shlib/shlib.bash && git_files "misc/dbt-materialize/*.py" | xargs bin/pyactivate -m ruff --target-version=py38{fix}',
150    ]
151
152
153if __name__ == "__main__":
154    os.chdir(MZ_ROOT)
155    exit(main())