dock

#! /usr/bin/python
"""
A wrapper script around docker.

Author: Huizi Mao
"""
from __future__ import print_function

import os
import subprocess
import argparse


def load_default_mapping():
    """Load directory mapping from dockpath config.
       
    Environment variables and symlinks will be resolved in dockpath.
    """
    filename = os.getenv("HOME") + '/.docker/dockpath'
    if os.path.exists(filename):
        lines = open(filename).readlines()
        mapping = []
        for line in lines:
            line = line.strip()
            if line.startswith('#') or len(line) == 0:
                continue
            if len(line.split(':')) != 2:
                raise ValueError("Cannot parse dockpath, error line: " + line)
            mapping.append(line.split(':'))
    else:
        mapping = []
    for idx in range(len(mapping)):
        mapp = mapping[idx]
        # Resolve environment variables
        mapp = (os.path.expandvars(mapp[0]),
                        os.path.expandvars(mapp[1]))

        # Resolve symlink
        if os.path.realpath(mapp[0]) != mapp[0]:
            mapp = (os.path.realpath(mapp[0]), mapp[1])
        mapping[idx] = mapp
    return mapping


def map_dir(mapping, dirname, default='/'):
    """Replace a file/directory with correpsonding path in the container.
    
    Only replace dirname that starts with /
    """
    if dirname.startswith('/'):
        for mapp in mapping:
            if dirname.startswith(mapp[0]):
                dirname = os.path.join(mapp[1], os.path.relpath(dirname, mapp[0]))
                return dirname
    return default


def shell_run(cmd):
    """Ensure the correct alias is called, e.g., nvidia-docker."""
    subprocess.call([os.environ['SHELL'], '-i', '-c', cmd])


def get_start_cmd(args, default_mapping):
    if args.x11:
        x11_mount = "-e DISPLAY=$DISPLAY -v /tmp/.X11-unix:/tmp/.X11-unix"
    else:
        x11_mount = ""
    dir_mount = ' '.join(['-v %s:%s'%(mapp[0], mapp[1]) for mapp in default_mapping])
    cmd = "docker run -ti {mount_cmd} {x11_mount} --ipc=host --name={name} {image} bash; docker start {name}"
    cmd = cmd.format(image=args.image, name=args.name, mount_cmd=dir_mount, x11_mount=x11_mount)
    hints = "Wait for bash prompt and then exit to finish setup."
    return cmd, hints


def get_run_cmd(args, default_mapping):
    workdir = map_dir(default_mapping, os.getcwd())
    home = map_dir(default_mapping, os.getenv('HOME'))

    # Every arg will be checked if starting with /.
    # If yes, it is considered to be a file or directory, and will
    # be mapped to container directory.
    mapped_args = [map_dir(default_mapping, arg, default=arg) for arg in args.cmd]
    if args.root:
        uid = 0
    else:
        uid = os.getuid()
    cmd = "docker exec -it -u {uid} -e HOME={home} -w {workdir} {name} sh -c '{command}'"
    cmd = cmd.format(workdir=workdir, name=args.name, command=' '.join(mapped_args), uid=uid, home=home)
    return cmd, ""


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--print', action='store_true', help='Print the custom docker comand only'
                                                             ' do not run')
    subparsers = parser.add_subparsers(dest='command')
    start_parser = subparsers.add_parser('start', help='Create and start a new container')
    start_parser.add_argument('name', help='Tag for the new container')
    start_parser.add_argument('image', help='Docker image to start')
    start_parser.add_argument('--x11', action='store_true', help='Add x11 forwarding')

    run_parser = subparsers.add_parser('run', help='Run command in an existing container')
    run_parser.add_argument('name', help='Tag of the existing container')
    run_parser.add_argument('--root', action='store_true', help='Run as root')
    run_parser.add_argument('cmd', nargs=argparse.REMAINDER, help='Custom command to run')

    args = parser.parse_args()

    default_mapping = load_default_mapping()

    if args.command == 'start':
        cmd, hints = get_start_cmd(args, default_mapping)
    elif args.command == 'run':
        cmd, hints = get_run_cmd(args, default_mapping)
    else:
        raise NotImplementedError

    if not args.print:
        print("Executing the following command:")
        print(cmd)
        print(hints)
        print("==============================")
        shell_run(cmd)
    else:
        print(cmd)


if __name__ == "__main__":
    main()