Write a dummy file to disk, with optional parameters. Usage: ./write_dummy_data.py <num_mb filename use_mib verbose byte_to_write max_chunk_size>

write_dummy_data.py

#!/usr/bin/env python3

from inspect import signature
import os
import sys

# TODO: fix the megabytes/mebibytes inconsistency in comments/docstrings.


def write_data(
    num_megabytes: int = 100,  # Write this many mebibytes (or megabytes)
    file: str = "data.bin",  # Filename
    use_mib: bool = True,  # If True, use mebibytes (MiB) instead of megabytes (MB)
    verbose: bool = True,  # Print debug info
    byte_to_write: bytes = b"0",  # Which byte character to fill the file with
    max_chunk_size_mb: int = 100,  # Flush to disk when the buffer reaches this size
) -> None:
    """Write "num_megabytes" megabytes of "byte_to_write" byte
    into "file" file, at the script folder."""
    assert len(byte_to_write) == 1, "Must be exactly 1 ASCII literal"
    assert max_chunk_size_mb >= 1
    full_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), file)
    assert not os.path.exists(full_path), f"Path already exists: {full_path}"
    unit = pow((1024 if use_mib else 1000), 2)
    acceptable_chunk_size_bytes = max_chunk_size_mb * unit
    data_total_size_in_bytes = num_megabytes * unit
    bytes_written_total = 0
    with open(full_path, mode="ab") as f:
        chunk_size = min(data_total_size_in_bytes, acceptable_chunk_size_bytes)
        while bytes_written_total < data_total_size_in_bytes:
            if bytes_written_total + chunk_size <= data_total_size_in_bytes:
                n_bytes = chunk_size
            else:
                n_bytes = data_total_size_in_bytes - bytes_written_total
            f.write(n_bytes * byte_to_write)
            bytes_written_total += n_bytes
    if verbose:
        print(
            f"Wrote {bytes_written_total} bytes "
            f'({bytes_written_total / unit} {"MiB" if use_mib else "MB"}) to path: "{full_path}"'
        )


def boolify(val):
    try:
        if val.lower() == "false" or val == "0":
            return False
    except AttributeError:
        pass
    return bool(val)


if __name__ == "__main__":
    args = {
        k: v
        for k, v in dict(
            zip(
                tuple(signature(write_data).parameters),
                (
                    *sys.argv[1:],
                    *(
                        [None]
                        * (
                            len(signature(write_data).parameters)
                            - len(sys.argv[1:])
                        )
                    ),
                ),
            )
        ).items()
        if v is not None
    }
    if "num_megabytes" in args:
        args["num_megabytes"] = int(args["num_megabytes"])
    if "use_mib" in args:
        args["use_mib"] = boolify(args["use_mib"])
    if "verbose" in args:
        args["verbose"] = boolify(args["use_mib"])
    if "byte_to_write" in args:
        args["byte_to_write"] = bytes(args["byte_to_write"].encode("utf-8"))
    if "max_chunk_size_mb" in args:
        args["max_chunk_size_mb"] = int(args["max_chunk_size_mb"])
    write_data(**args)