path: root/python/refcat/cli.py

"""
              ____           __
   ________  / __/________ _/ /_
  / ___/ _ \/ /_/ ___/ __ `/ __/
 / /  /  __/ __/ /__/ /_/ / /_
/_/   \___/_/  \___/\__,_/\__/

Command line entry point for running various data tasks.

    $ refcat.pyz TASK [OPTIONS]

    $ refcat.pyz COMMAND [OPTIONS]

Commands: ls, ll, deps, tasks, files, config, cat, completion

To install completion run:

    $ source <(refcat.pyz completion)
"""

import logging
import os
import subprocess
import sys
import tempfile

import luigi
from luigi.cmdline_parser import CmdlineParser
from luigi.parameter import MissingParameterException
from luigi.task import Register
from luigi.task_register import TaskClassNotFoundException

from refcat import __version__
from refcat.deps import dump_deps, dump_deps_dot
from refcat.settings import LOGGING_CONF_FILE, settings
from refcat.tasks import *
from refcat.techreport import *
from refcat.utils import columnize

# These are utility classes of luigi.
suppress_task_names = [
    "Available",
    "BaseTask",
    "Config",
    "Executable",
    "ExternalTask",
    "RangeBase",
    "RangeByMinutes",
    "RangeByMinutesBase",
    "RangeDaily",
    "RangeDailyBase",
    "RangeHourly",
    "RangeHourlyBase",
    "RangeMonthly",
    "Task",
    "TestNotificationsTask",
    "WrapperTask",
]


def effective_task_names():
    """
    Runnable, relevant task names.
    """
    names = (name for name in sorted(Register.task_names()))
    names = (name for name in names if name not in suppress_task_names and not name.islower())
    return names


def tasks():
    """
    Print task names.
    """
    for name in effective_task_names():
        print(name)


def files():
    """
    Print task names and output file.
    """
    for name in effective_task_names():
        klass = Register.get_task_cls(name)
        try:
            print("{:40s}{}".format(name, klass().output().path))
        except AttributeError:
            print("{:40s}".format(name))


def cat():
    """
    File contents to stdout.
    """
    if len(sys.argv) < 2:
        raise ValueError("task name required")
    try:
        parser = CmdlineParser(sys.argv[2:])
        output = parser.get_task_obj().output()
        filename = output.path
        if filename.endswith(".zst"):
            subprocess.run(["zstdcat", "-T0", filename])
        elif filename.endswith(".gz"):
            subprocess.run(["zcat", filename])
        else:
            subprocess.run(["cat", filename])
    except FileNotFoundError:
        print("file not found: {}".format(filename), file=sys.stderr)
    except AttributeError:
        print("most likely not a task object", file=sys.stderr)
    except TaskClassNotFoundException as exc:
        print("no such task")


def ls():
    """
    Print output filename for task.
    """
    if len(sys.argv) < 2:
        raise ValueError("task name required")
    try:
        parser = CmdlineParser(sys.argv[2:])
        output = parser.get_task_obj().output()
    except TaskClassNotFoundException as exc:
        print("no such task")
    else:
        print(output.path)


def ll():
    """
    Long output.
    """
    if len(sys.argv) < 2:
        raise ValueError("task name required")
    try:
        parser = CmdlineParser(sys.argv[2:])
        output = parser.get_task_obj().output()
        filename = output.path
        subprocess.run(["ls", "-lah", filename])
    except FileNotFoundError:
        print("file not found: {}".format(filename), file=sys.stderr)
    except AttributeError:
        print("most likely not a task object", file=sys.stderr)
    except TaskClassNotFoundException as exc:
        print("no such task")


def deps():
    """
    Render task dependencies.
    """
    if len(sys.argv) < 2:
        raise ValueError("task name required")
    try:
        parser = CmdlineParser(sys.argv[2:])
        obj = parser.get_task_obj()
        dump_deps(obj)
    except TaskClassNotFoundException as exc:
        print("no such task")


def dot():
    """
    Render task dependencies as dot.
    """
    if len(sys.argv) < 2:
        raise ValueError("task name required")
    try:
        parser = CmdlineParser(sys.argv[2:])
        obj = parser.get_task_obj()
        dump_deps_dot(obj)
    except TaskClassNotFoundException as exc:
        print("no such task")


def config():
    """
    Dump config to stdout.
    """
    with open(settings.settings_file) as f:
        print(f.read())


def run():
    """
    For uniformity, have an extra subcommand for running a task.
    """
    raise NotImplementedError()


def completion():
    """
    TODO: completion snippet can live elsewhere.
    """
    snippet = """
_refcat_completion()
{
    local cur prev
    COMPREPLY=()
    cur=${COMP_WORDS[COMP_CWORD]}
    prev=${COMP_WORDS[COMP_CWORD-1]}

    if [ $COMP_CWORD -eq 1 ]; then
        COMPREPLY=( $(compgen -W "cat completion config deps files ll ls tasks" -- $cur) )
    elif [ $COMP_CWORD -eq 2 ]; then
        case "$prev" in
            "cat")
                COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) )
                ;;
            "ll")
                COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) )
                ;;
            "ls")
                COMPREPLY=( $(compgen -W "$(refcat.pyz tasks)" -- $cur) )
                ;;
            *)
                ;;
        esac
    fi
    return 0
}

complete -F _refcat_completion "refcat.pyz"
complete -F _refcat_completion "refcat"
    """
    print(snippet)


def main():
    """
    The main command line entry point.
    """
    if os.path.exists(LOGGING_CONF_FILE):
        logging.config.fileConfig(LOGGING_CONF_FILE)

    if settings.TMPDIR:
        # This might not be passed on to subprocesses.
        tempfile.tempdir = settings.TMPDIR

    # Explicitly name valid subcommands.
    sub_commands = [
        "cat",
        "completion",
        "config",
        "deps",
        "dot",
        "files",
        "ll",
        "ls",
        "tasks",
    ]

    if len(sys.argv) >= 2 and sys.argv[1] in sub_commands:
        try:
            # A hack to call the function within this file matching the
            # subcommand name.
            func = globals()[sys.argv[1]]
            if callable(func):
                func()
            else:
                print("not implemented: {}".format(sys.argv[1]))
        except KeyError:
            print("subcommand not implemented: {}".format(sys.argv[1]), file=sys.stderr)
        except (AttributeError, ValueError, RuntimeError) as exc:
            print(exc, file=sys.stderr)
            sys.exit(1)
        else:
            sys.exit(0)
    elif len(sys.argv) == 1:
        shiv_root_default = os.path.join(os.path.expanduser("~"), ".shiv")
        print(__doc__)
        print("VERSION   {}".format(__version__))
        print("SETTINGS  {}".format(settings.settings_file))
        print("BASE      {}".format(settings.BASE))
        print("TAG       {}".format(os.path.join(settings.BASE, Refcat.TAG)))
        print("TMPDIR    {}".format(settings.TMPDIR))
        print("SHIV_ROOT {}".format(os.environ.get("SHIV_ROOT") or shiv_root_default))
        print()
        print(columnize(list(effective_task_names())))
        sys.exit(0)

    # If we found no subcommand, assume task name.
    try:
        luigi.run(local_scheduler=True)  # XXX: add logging_conf_file
    except MissingParameterException as err:
        print('missing parameter: %s' % err, file=sys.stderr)
        sys.exit(1)
    except TaskClassNotFoundException as err:
        print(err, file=sys.stderr)
        sys.exit(1)