multinetspeeder/scripts/rotate-log.py

#!/usr/bin/env python3
from __future__ import annotations

import os
import sys
import time
from pathlib import Path


def rotate(path: Path, backups: int) -> None:
    if backups <= 0:
        path.write_text("")
        return
    oldest = path.with_name(f"{path.name}.{backups}")
    if oldest.exists():
        oldest.unlink()
    for index in range(backups - 1, 0, -1):
        source = path.with_name(f"{path.name}.{index}")
        target = path.with_name(f"{path.name}.{index + 1}")
        if source.exists():
            source.replace(target)
    if path.exists():
        path.replace(path.with_name(f"{path.name}.1"))


def main() -> int:
    if len(sys.argv) != 4:
        print("usage: rotate-log.py <log_path> <max_bytes> <backups>", file=sys.stderr)
        return 1
    log_path = Path(sys.argv[1])
    max_bytes = int(sys.argv[2])
    backups = int(sys.argv[3])
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_path.touch(exist_ok=True)

    stream = sys.stdin.buffer
    while True:
        chunk = stream.readline()
        if not chunk:
            break
        chunk = f"[{int(time.time() * 1000)}] ".encode() + chunk
        with log_path.open("ab") as handle:
            handle.write(chunk)
        try:
            if log_path.stat().st_size > max_bytes:
                rotate(log_path, backups)
                log_path.touch(exist_ok=True)
        except FileNotFoundError:
            log_path.touch(exist_ok=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())