Rewrite tarfiles so that they are folder-prefixed

make_folder_prefixed.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  make_folder_prefixed.py
#

import os
import tarfile
import tempfile
import shutil
import contextlib

from pathlib import Path

import click
import sh


@click.command()
@click.argument("input_file")
@click.option("--fix", is_flag=True, default=False)
def main(input_file: str, fix=False):
    """
    Check if a .tar.bz2 file was compressed with a self-contained folder, i.e.
    that when foo.tar.bz2 decompresses to foo/. By default, just audits the file,
    but with --fix it will rewrite it to fix this issue.
    """
    assert input_file.endswith(".tar.bz2")

    stem = get_stem(input_file)
    if is_tarfile_prefixed(input_file, stem):
        print("GOOD", input_file)
    elif fix:
        print("FIX ", input_file)
        rewrite_with_stem(input_file, stem)
    else:
        print("BAD ", input_file)


def is_tarfile_prefixed(input_file: str, stem: str) -> bool:
    with tarfile.open(input_file) as t:
        return all(info.name.startswith(stem) for info in t)


def rewrite_with_stem(input_file: str, stem: str) -> None:
    input_file = os.path.abspath(input_file)

    with tempfile.TemporaryDirectory() as d:
        with chdir(d):
            os.mkdir(stem)
            with chdir(stem):
                run(f"tar xf {input_file}")

            run(f"tar c {stem} | pbzip2 -c >{stem}.tar.bz2")
            shutil.move(f"{stem}.tar.bz2", input_file)


def run(command: str) -> None:
    code = os.system(command)
    if code != 0:
        raise Exception(f"non-zero exit code: {code}")


def get_stem(input_file: str) -> str:
    return Path(input_file).name[: -len(".tar.bz2")]


@contextlib.contextmanager
def chdir(d):
    current_dir = os.getcwd()
    try:
        os.chdir(d)
        yield
    finally:
        os.chdir(current_dir)


if __name__ == "__main__":
    main()