tempwrap

#!/usr/bin/env python3

# This script is a generic wrapper that copies a directory (or
# parts of it) into a temporary location (under $TMP), executes
# a command from that location, and copies the results back.
#
# Such a wrapper is useful to force all temporary files of a
# command to reside on a fast local disk.

import fnmatch
import getopt
import glob
import os
import re
import subprocess
import sys
import tempfile
import time

def usage():
    print("usage: tempwrap [-i pattern] [-o pattern] [-e sed-arg] ... fromdir todir -- command [arg ...]")
    exit(1)

opts, args = getopt.gnu_getopt(sys.argv[1:], "i:o:e:")

i_patterns = []
o_patterns = []
sed_args = []
for opt, arg in opts:
    if opt == "-i": i_patterns.append(arg)
    if opt == "-o": o_patterns.append(arg)
    if opt == "-e": sed_args.append("-e"); sed_args.append(arg)
if len(args) < 3:
    usage()
srcdir = args[0]
dstdir = args[1]
command = args[2:]

def mkdir(path):
    os.makedirs(path, exist_ok=True)

def copy(srcpath, dstpath, sed_args):
    if sed_args:
        with open(dstpath, "wb") as f:
            if srcpath.endswith(".gz"):
                p1 = subprocess.Popen(["zcat", srcpath], stdout=subprocess.PIPE)
                p2 = subprocess.Popen(["sed", *sed_args], stdin=p1.stdout, stdout=subprocess.PIPE)
                p3 = subprocess.Popen(["gzip", "-c"], stdin=p2.stdout, stdout=f)
                if p3.wait() != 0: raise ValueError("gzip failed")
                if p2.wait() != 0: raise ValueError("sed failed")
                if p1.wait() != 0: raise ValueError("zcat failed")
            else:
                subprocess.check_call(["sed", *sed_args, srcpath], stdout=f)
    else:
        subprocess.check_call(["cp", "-aL", srcpath, dstpath])

def copytree(srcdir, dstdir, patterns, sed_args):
    print("* copy from", repr(srcdir), "into", repr(dstdir))
    rx = re.compile("|".join(fnmatch.translate(p) for p in patterns)) if patterns else None
    done_dirs = set()
    for dirname, dirs, files in os.walk(srcdir):
        for dir in sorted(dirs):
            fullpath = os.path.join(dirname, dir)
            relpath = os.path.relpath(fullpath, srcdir)
            if rx is None or rx.match(relpath):
                if relpath not in done_dirs:
                    print("* mkdir", repr(relpath))
                    mkdir(os.path.join(dstdir, relpath))
                    done_dirs.add(relpath)
        for file in sorted(files):
            fullpath = os.path.join(dirname, file)
            relpath = os.path.relpath(fullpath, srcdir)
            if rx is None or rx.match(relpath):
                reldirname = os.path.dirname(relpath)
                if reldirname not in done_dirs:
                    print("* mkdir", repr(reldirname))
                    mkdir(os.path.join(dstdir, reldirname))
                    done_dirs.add(reldirname)
                print("* copy", repr(relpath))
                copy(fullpath, os.path.join(dstdir, relpath), sed_args)

t0 = time.time()
with tempfile.TemporaryDirectory(prefix="wrap-") as tmpdir:
    copytree(srcdir, tmpdir, i_patterns, sed_args)
    print("* run", *command)
    sys.stdout.flush()
    t1 = time.time()
    subprocess.check_call(command, cwd=tmpdir)
    t2 = time.time()
    copytree(tmpdir, dstdir, o_patterns, [])
t3 = time.time()
print(f"* copy in time: {t1-t0:.4g}s")
print(f"* copy out time: {t3-t2:.4g}s")
print(f"* run time: {t2-t1:.4g}s")