A simple CLI to copy directory trees and keep trying upon errors.

resilient_tree_copy.py

#!/usr/bin/env python3

# Copyright (c) 2020, Moritz E. Beber.


"""Copy entire directory trees and keep trying on errors."""


import argparse
import logging
import os
import pickle
from pathlib import Path
from shutil import copy2
from typing import Collection

from tenacity import retry, stop_after_attempt, wait_fixed, before_sleep_log


logger = logging.getLogger()


def parse_args():
    parser = argparse.ArgumentParser(description="Find the latest Docker image tag.")
    parser.add_argument(
        "source",
        metavar="SOURCE",
        help="The root of the source directory tree to copy.",
    )
    parser.add_argument(
        "destination",
        metavar="DESTINATION",
        help="The destination where to copy the files to.",
    )
    parser.add_argument(
        "--verbosity",
        help="The desired log level (default WARNING).",
        choices=("CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"),
        default="WARNING",
    )
    args = parser.parse_args()
    logging.basicConfig(level=args.verbosity, format="[%(levelname)s] %(message)s")
    return args


@retry(
    stop=stop_after_attempt(10),
    wait=wait_fixed(0.1),
    before_sleep=before_sleep_log(logger, logging.ERROR),
)
def build_tree(source: Path, cache: Path) -> Collection[Path]:
    """
    Collect all file paths from a directory tree.

    Since walking the directory tree and copying may fail, we assemble the tree first
    so that we can resume copying later.

    """
    if cache.exists():
        with cache.open("rb", encoding=None) as handle:
            tree = pickle.load(handle)
    else:
        tree = set()
    for dir_path, dir_names, filenames in os.walk(str(source)):
        for name in filenames:
            path = Path(dir_path) / name
            logger.debug("%r", str(path))
            tree.add(path)
    with cache.open("wb", encoding=None) as handle:
        pickle.dump(tree, handle)
    return tree


@retry(
    stop=stop_after_attempt(10),
    wait=wait_fixed(0.1),
    before_sleep=before_sleep_log(logger, logging.ERROR),
)
def copy_file(source: Path, destination: Path, path: Path) -> None:
    """"""
    relative_path = path.relative_to(source)
    target = destination / relative_path
    target.parent.mkdir(parents=True, exist_ok=True)
    copy2(str(path), str(target))


def copy_tree(source: Path, destination: Path, tree: Collection[Path]) -> None:
    """"""
    for path in tree:
        copy_file(source, destination, path)


def main() -> None:
    """"""
    args = parse_args()
    source = Path(args.source)
    destination = Path(args.destination)
    cache = destination / ".tree_cache.pkl"
    tree = build_tree(source, cache)
    copy_tree(source, destination, tree)
    cache.unlink()


if __name__ == "__main__":
    main()