like `split`, but with better numerical naming

split.py

#!/usr/bin/env python
import argparse
from pathlib import Path


def split_file(filename: str | Path, lines_per_file: int, output_prefix: str):
    """Split a text file into multiple files.
    
    Parameters
    ----------
    filename : str or Path
        The name of the input file.
    lines_per_file : int
        The maximum number of lines per output file.
    output_prefix : str
        The prefix to use for the output file names.


    Examples
    --------
    >>> split_file("input.txt", 1000, "output")

    This will split the file "input.txt" into multiple files with names like
    "output_0001.txt", "output_0002.txt", etc. Each output file will contain
    1000 lines from the input file, except for the last file, which may have
    fewer lines if the total number of lines in the input file is not a
    multiple of 1000.
    """
    # Count total number of lines in the file to get the number of digits for zero-padding
    with open(filename, "r") as f:
        total_lines = sum(1 for line in f)

    num_files = -(-total_lines // lines_per_file)  # Ceiling division
    num_digits = len(str(num_files))

    with open(filename, "r") as f:
        current_line = 0
        current_file = 0
        current_out_file = None

        for line in f:
            if current_line % lines_per_file == 0:
                if current_out_file:
                    current_out_file.close()

                current_file += 1
                suffix = str(current_file).zfill(num_digits)
                current_out_file = open(f"{output_prefix}_{suffix}.txt", "w")

            current_out_file.write(line)
            current_line += 1

        if current_out_file:
            current_out_file.close()


def main():
    parser = argparse.ArgumentParser(
        description="Split a text file into multiple files."
    )
    parser.add_argument("filename", type=str, help="The name of the input file.")
    parser.add_argument(
        "lines_per_file", type=int, help="The maximum number of lines per output file."
    )
    parser.add_argument(
        "output_prefix", type=str, help="The prefix to use for the output file names."
    )
    args = parser.parse_args()

    split_file(args.filename, args.lines_per_file, args.output_prefix)


if __name__ == "__main__":
    main()