urd.py

#!/usr/bin/env python3

''' upload - run - download '''

from pathlib import Path
import os
import string
import subprocess
import sys
import shlex

# Make GNU rsync >= 3.2.4 treat remote paths like older rsync and openrsync do
# (word-split by the remote shell), so our shlex.quote() works the same everywhere.
RSYNC_ENV = {**os.environ, 'RSYNC_OLD_ARGS': '1'}


def main():
    [machine, *args] = sys.argv[1:]
    dirname = Path.cwd().name

    print('*' * 20, 'uploading ...')
    subprocess.check_call([
        'rsync',
        '-avzc',
        '--filter', ':- .gitignore',
        '.',
        machine + ':' + shlex.quote(dirname),
    ], env=RSYNC_ENV)

    if args:
        args = [shlex.quote(a) for a in args]
        print('*' * 20, 'running', ' '.join(args), '...')
        p = subprocess.run([
            'ssh',
            '-t',
            machine,
            'cd ' + shlex.quote(dirname) + ' && bash -l -i -c ' + shlex.quote(' '.join(args)),
        ])
        if p.returncode:
            print('*' * 20, 'remote command failed with exit code', p.returncode)
            sys.exit(1)

    d = Path('.urd_download')
    if d.exists():
        print('*' * 20, 'downloading ...')
        for line in d.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            src, arrow, dst = line.partition('->')
            src = string.Template(src.strip()).substitute(machine=machine)
            dst = string.Template(dst.strip()).substitute(machine=machine) if arrow else src
            Path(dst).parent.mkdir(parents=True, exist_ok=True)
            # Failures (e.g. optional outputs missing on the remote) are not fatal;
            # rsync prints its own warnings.
            subprocess.run([
                'rsync',
                '-avzL',
                machine + ':' + shlex.quote(dirname + '/' + src),
                dst,
            ], env=RSYNC_ENV)


if __name__ == '__main__':
    main()