Skip to content

Instantly share code, notes, and snippets.

@Midnighter
Created May 19, 2024 06:39
Show Gist options
  • Save Midnighter/54e165459a4209639e3ab732cea84087 to your computer and use it in GitHub Desktop.
Save Midnighter/54e165459a4209639e3ab732cea84087 to your computer and use it in GitHub Desktop.
A simple CLI to copy directory trees and keep trying upon errors.
#!/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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment