holes

#!/usr/bin/env python3

import io
import math
import optparse
import os
import sys
from dataclasses import dataclass
from typing import Optional


NULL_CHAR = '_'
FILL_CHAR = 'X'

MAX_BLOCK_SIZE = 100 * 1024 * 1024


class HolesError(Exception):
    pass


def sizeof_fmt(num, suffix='B') -> str:
    if num == 0:
        return "0B"
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def main() -> int:
    usage = '''%prog [options] FILE...

Find runs of null bytes in files, print the percentage of blocks that are
filled with content vs empty. By default, the blocks will be dynamically sized
based on the file's total size and the width of the output terminal. Pass -b to
use a fixed block size instead.
'''.rstrip()
    p = optparse.OptionParser(usage=usage)
    p.add_option('-q', '--quiet', dest='verbose', action='store_false',
                 help="Silence visualization", default=True)
    p.add_option('-b', '--blocksize', dest='blocksize', metavar='BYTES',
                 type='int', help='Size of blocks to examine')

    opts, args = p.parse_args()

    filenames = args

    if not filenames:
        p.print_help()
        return 1

    try:
        for filename in filenames:
            print_file_info(
                filename=filename,
                print_dots=opts.verbose,
                blocksize=opts.blocksize,
                verbose=opts.verbose,
            )
    except HolesError as err:
        sys.stderr.write(f"\n\terror: {err}\n")
        return 2

    return 0


def print_file_info(
    filename: str,
    blocksize: Optional[int] = None,
    print_dots: bool = True,
    verbose: bool = False,
) -> None:
    sys.stdout.write(filename)

    res = count_filled(
        filename=filename,
        blocksize=blocksize,
        print_dots=print_dots,
    )

    # technically / 0 is NaN, but 0% makes intuitive sense
    frac = res.count / res.total if res.total > 0 else 0
    out = f"\t{frac:.0%} full of {sizeof_fmt(res.file_size)}"

    print(out)


@dataclass
class CountResult:
    filename: str
    count: int
    total: int
    file_size: int
    block_size: int


def count_filled(
    filename: str,
    blocksize: Optional[int] = None,
    print_dots: bool = True,
    print_blocksize: bool = True,
) -> CountResult:
    count = 0
    total = 0

    fh = io.FileIO(filename, 'rb')

    # Seek to end to determine file size.
    # This method works on block devices, unlike stat.
    file_size = fh.seek(0, os.SEEK_END)
    fh.seek(0)

    if blocksize is None:
        if sys.stdout.isatty():
            max_blocks = os.get_terminal_size().columns - 9
        else:
            max_blocks = 71

        if max_blocks < 1:
            max_blocks = 1

        blocksize = math.ceil(file_size / max_blocks)
        blocksize = min(blocksize, MAX_BLOCK_SIZE)

    null_comparison = memoryview(b'\0' * blocksize)

    buf = bytearray(blocksize)

    if print_blocksize:
        if print_dots:
            sys.stdout.write("\n\t")
        else:
            sys.stdout.write("\t")
        sys.stdout.write("bs=" + sizeof_fmt(blocksize))

    if print_dots:
        sys.stdout.write("\n\t")

    sys.stdout.flush()

    while True:
        try:
            size = fh.readinto(buf)
        except OSError as err:
            pos = fh.tell()
            raise HolesError(f"{err} at fd offset={pos}") from err

        if size == 0:
            break

        total += 1

        if buf.startswith(null_comparison[:size], 0, size):
            if print_dots:
                sys.stdout.write(NULL_CHAR)
                sys.stdout.flush()
        else:
            count += 1
            if print_dots:
                sys.stdout.write(FILL_CHAR)
                sys.stdout.flush()

    if print_dots:
        sys.stdout.write('\n')
        sys.stdout.flush()

    return CountResult(
        filename=filename,
        count=count,
        total=total,
        file_size=file_size,
        block_size=blocksize,
    )


if __name__ == '__main__':
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)