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
 22def main() -> int:
 23    parser = argparse.ArgumentParser(prog="fmt")
 24    parser.add_argument("--check", action="store_true")
 25    args = parser.parse_args()
 26
 27    tasks: list[tuple[str, TaskSpec]] = [
 28        ("rustfmt", _rustfmt_fn(check=args.check)),
 29        ("buf", _buf_cmd(check=args.check)),
 30    ]
 31
 32    if not os.environ.get("MZDEV_NO_PYTHON"):
 33        tasks += [
 34            ("black", _black_cmd(check=args.check)),
 35            ("ruff", _ruff_cmd(check=args.check)),
 36            ("ruff-dbt", _ruff_dbt_cmd(check=args.check)),
 37        ]
 38
 39    return 1 if run_parallel(tasks, spinner_suffix="formatters") else 0
 40
 41
 42def _rustfmt_fn(*, check: bool):
 43    """Return a callable that runs cargo metadata + parallel rustfmt."""
 44
 45    def run() -> tuple[bool, str]:
 46        ncpus = os.cpu_count() or 8
 47
 48        result = subprocess.run(
 49            ["cargo", "metadata", "--no-deps", "--format-version=1"],
 50            capture_output=True,
 51            text=True,
 52        )
 53        if result.returncode != 0:
 54            return False, result.stderr.strip()
 55
 56        meta = json.loads(result.stdout)
 57        kinds = {"lib", "bin", "bench", "test", "example", "proc-macro", "custom-build"}
 58        paths = [
 59            t["src_path"]
 60            for pkg in meta["packages"]
 61            for t in pkg["targets"]
 62            if kinds & set(t["kind"])
 63        ]
 64        if not paths:
 65            return True, ""
 66
 67        # Split into batches and run rustfmt in parallel.
 68        batch_size = math.ceil(len(paths) / ncpus)
 69        batches = [paths[i : i + batch_size] for i in range(0, len(paths), batch_size)]
 70
 71        cmd_base = ["rustfmt", "--config", "error_on_line_overflow=true"]
 72        if check:
 73            cmd_base.append("--check")
 74
 75        procs = [
 76            subprocess.Popen(
 77                cmd_base + batch,
 78                stdout=subprocess.PIPE,
 79                stderr=subprocess.STDOUT,
 80            )
 81            for batch in batches
 82        ]
 83
 84        all_output = []
 85        all_ok = True
 86        for proc in procs:
 87            stdout, _ = proc.communicate()
 88            if proc.returncode != 0:
 89                all_ok = False
 90            out = stdout.decode("utf-8").strip()
 91            if out:
 92                all_output.append(out)
 93
 94        return all_ok, "\n".join(all_output)
 95
 96    return run
 97
 98
 99def _buf_cmd(*, check: bool) -> list[str]:
100    if check:
101        return ["buf", "format", "src", "--diff", "--exit-code"]
102    return ["buf", "format", "src", "-w"]
103
104
105def _black_cmd(*, check: bool) -> list[str]:
106    args = "--check --quiet" if check else "--quiet"
107    return [
108        "bash",
109        "-c",
110        f'. misc/shlib/shlib.bash && git_files "*.py" | xargs bin/pyactivate -m black {args}',
111    ]
112
113
114def _ruff_cmd(*, check: bool) -> list[str]:
115    fix = "" if check else " --fix"
116    return [
117        "bash",
118        "-c",
119        f'. misc/shlib/shlib.bash && git_files "*.py" | grep -v "^misc/dbt-materialize/" | xargs bin/pyactivate -m ruff{fix}',
120    ]
121
122
123def _ruff_dbt_cmd(*, check: bool) -> list[str]:
124    fix = "" if check else " --fix"
125    return [
126        "bash",
127        "-c",
128        f'. misc/shlib/shlib.bash && git_files "misc/dbt-materialize/*.py" | xargs bin/pyactivate -m ruff --target-version=py38{fix}',
129    ]
130
131
132if __name__ == "__main__":
133    os.chdir(MZ_ROOT)
134    exit(main())